Tensorflow Keras Model subclassing -- call function

I am experimenting with self supervised learning using tensorflow. The example code I'm running can be found in the Keras examples website. This is the link to the NNCLR example. The Github link to download the code can be found here. While I have no issues running the examples, I am running into issues when I try to save the pretrained or the finetuned model using model.save().
The error I'm getting is this:

 f"Model {model} cannot be saved either because the input shape is not "
 ValueError: Model <__main__.NNCLR object at 0x7f6bc0f39550> cannot be saved either 
 because the input shape is not available or because the forward pass of the model is
 not defined. To define a forward pass, please override `Model.call()`. 
 To specify an input shape, either call `build(input_shape)` directly, or call the model on actual data using `Model()`, `Model.fit()`, or `Model.predict()`. 
 If you have a custom training step, please make sure to invoke the forward pass in train step through  
`Model.__call__`, i.e. `model(inputs)`, as opposed to `model.call()`.

I am unsure how to override the Model.call() method. Appreciate some help.

One way to achieve model saving in such cases is to override the save (or save_weights) method in the keras.Model class. In your case, first initialize the finetune model in the NNCLR class. And next, override the save method for it. FYI, in this way, you may also able to use ModelCheckpoint API.

As said, define the finetune model in the NNCLR model class and override the save method for it.

class NNCLR(keras.Model):
    def __init__(...):
        super().__init__()
        ...
        self.finetuning_model = keras.Sequential(
            [
                layers.Input(shape=input_shape),
                self.classification_augmenter,
                self.encoder,
                layers.Dense(10),
            ],
            name="finetuning_model",
        )
    
    ...

    def save(
        self, filepath, overwrite=True, include_optimizer=True, 
        save_format=None, signatures=None, options=None
    ):
        self.finetuning_model.save(
            filepath=filepath,
            overwrite=overwrite,
            save_format=save_format,
            options=options,
            include_optimizer=include_optimizer,
            signatures=signatures
        )

model = NNCLR(...)
model.compile
model.fit

Next, you can do

model.save('finetune_model') # SavedModel format
finetune_model = tf.keras.models.load_model('finetune_model', compile=False)

'''
NNCLR code example: Evaluate sections
"A popular way to evaluate a SSL method in computer vision or 
for that fact any other pre-training method as such is to learn 
a linear classifier on the frozen features of the trained backbone 
model and evaluate the classifier on unseen images."
'''
for layer in finetune_model.layers:
    if not isinstance(layer, layers.Dense):
        layer.trainable = False

finetune_model.summary() # OK
finetune_model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")],
)
finetune_model.fit