Save tf model in colab to drive

Question

I have trained a tensorflow model for chatbot using Transformers in colab and now I am trying to save the model. I have mounted the drive and tried saving the model using the following.

from google.colab import drive
drive.mount('/content/gdrive')

model.save("Cornell_Chatbot")

But it keeps me throwing out this error.

AttributeError: 'PositionalEncoding' object has no attribute 'position'

Why this happens and how can I save my trained model?

EDIT:

Try 02 (TENSORFLOW CHECKPOINT)

checkpoint_path = "training_1/cp.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)

#Create checkpoint callback
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)
model.fit(dataset, epochs=EPOCHS,
          callbacks=[cp_callback])  

I have tried with tensofrflow checkpoint. I am using TPU to run the model. This gives me the following error:

UnimplementedError: File system scheme '[local]' not implemented (file: 'training_1/cp.ckpt_temp_44208607b0084dbcacbf799c6eb30c94/part-00000-of-00001')

I hope this occures since Cloud TPUs can only access data in GCS. Am I correct? How can I overcome this?