Fit model function in Tensorflow 2

fit-model.py

batch_size = 64

"""
Training the model for 60 epochs using our dataset.
The batch size (64) is the same for the validation data.
Only 1 callback was used, but could be more like TensorBoard, ModelCheckpoint, etc.
"""
history = model.fit(train_ds.batch(batch_size=batch_size),
                    epochs=60,
                    validation_data=validation_ds.batch(batch_size=batch_size),
                    callbacks=[early_callback])
model.save('model_base')