Sayak_Paul:

If your subclassed model (where I am overriding train_step()) contains two or more models and if you are passing ModelCheckpoint callback while calling .fit() on the subclassed model the callback would get confused.


That shouldn’t be like that. Models are supposed to be nestable.
…
The problem here is that the callback is defaulting to saving the model in HDF5 format (which apparently requires that to call .fit to set the input shape, and we don’t call fit on the nested mopdels.).
Set save_weights_only=True to save in the tensorflow checkpoint format and then it works.