TypeError: __init__() got an unexpected keyword argument 'reduction' while trying to load keras model - tensorflow

Question

I'm trying to load a Keras sequential model and get some error of unexpected keyword reduction.
Here's the code:
I've trained the following Sequential model:
model = Sequential([
tf.keras.Input(shape=(in_dim,)),
layers.Dense(
units=(in_dim+1),
activation=layers.LeakyReLU(alpha=.01)
),
layers.Dropout(rate=.05),
layers.Dense(
units=(in_dim),
activation=layers.LeakyReLU(alpha=.01)
),
layers.Dropout(rate=.05),
layers.Dense(units=1, activation="sigmoid")
])
and compile and fit it by Keras's tutorial:
def compile_and_fit(
model,
name,
X_y_train,
X_y_val,
optimizer=None,
max_epochs=10000,
batch_size=BATCH_SIZE,
):
X_train, y_train = X_y_train
steps_per_epoch = len(X_train) // batch_size
X_val, y_val = X_y_val
steps_per_epoch_val = len(X_val) // batch_size
if optimizer is None:
optimizer = get_optimizer(steps_per_epoch)
model.compile(
optimizer=optimizer,
loss=tf.keras.losses.BinaryCrossentropy(),
metrics=[
tf.keras.losses.BinaryCrossentropy(name="binary_crossentropy"),
tf.keras.metrics.Precision(name="precision"),
"accuracy",
],
)
model.summary()
history = model.fit(
x=X_train,
y=y_train,
steps_per_epoch=steps_per_epoch,
batch_size=batch_size,
epochs=max_epochs,
validation_data=X_y_val,
validation_steps=steps_per_epoch_val,
callbacks=get_callbacks(name),
verbose=1,
)
return history
And at the end saving the model:
model.save(f"./saved_model/my_model", save_format="tf")
But When I tried to load it:
model = tf.keras.models.load_model("./saved_model/my_model")
And got the following error:
TypeError Traceback (most recent call last)
Cell In [25], line 2
----> 2 model = tf.keras.models.load_model("./saved_model/my_model")
File /mnt/c/Code/venv/lib/python3.9/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
67 filtered_tb = _process_traceback_frames(e.__traceback__)
68 # To get the full stack trace, call:
69 # `tf.debugging.disable_traceback_filtering()`
---> 70 raise e.with_traceback(filtered_tb) from None
71 finally:
72 del filtered_tb
File /mnt/c/Code/venv/lib/python3.9/site-packages/keras/dtensor/utils.py:144, in inject_mesh.<locals>._wrap_function(instance, *args, **kwargs)
142 if mesh is not None:
143 instance._mesh = mesh
--> 144 init_method(instance, *args, **kwargs)
TypeError: __init__() got an unexpected keyword argument 'reduction'
Any ideas on how to fix it?