float and integer formatting in a pandas DataFrame

I have a data frame that I want to format in float for one part of the data frame and whole numbers for the other part. Basically, I want to format the data frame with three decimals if the decimals are different than zero and to whole numbers if the three decimals are zeros.

Data type
A    float64
B    float64
C    float64
D    float64
E    float64
dtype: object

I used the following code:

# Custom formatting function
def custom_format(value):
    if pd.isnull(value) or np.isinf(value):
        return value
    rounded_value = round(value, 3)
    if rounded_value == int(rounded_value):
        formatted_value = f'{int(rounded_value):.0f}'  # Show whole number only
    else:
        formatted_value = f'{rounded_value:.3f}'  # Show three decimals
    return formatted_value

# Apply custom formatting to the entire DataFrame
formatted_df = df.applymap(custom_format)

formatted_df

This is what I got:

This is what I want, however, the data type changed and become this:

A    object
B    object
C    object
D    object
E    object
dtype: object

The problem is that I need the values to be numerical while keeping the previous format. I tried multiple codes, but the object format always remains.

You can convert columns back to numerical type in the end of your code

formatted_df = formatted_df.apply(pd.to_numeric)