CancelledError: [_Derived_]RecvAsync is cancelled. Running on google colab

I am getting the following error when fitting my model. Occurs after every 5th epoch!

Link to notebook

Tensorflow, Keras, Cuda versions

Tried the following options:

1. TF_FORCE_GPU_ALLOW_GROWTH=1
2. Changing cuda and keras versions.

Error:

WARNING:tensorflow:From <ipython-input-21-44f42c6a15c1>:2: Model.fit_generator (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.
Instructions for updating:
Please use Model.fit, which supports generators.
Epoch 1/10
39/39 [==============================] - 10s 248ms/step - loss: 4.7112 - accuracy: 0.1422 - val_loss: 4.2005 - val_accuracy: 0.1601
Epoch 2/10
39/39 [==============================] - 8s 207ms/step - loss: 3.9892 - accuracy: 0.1991 - val_loss: 3.9835 - val_accuracy: 0.1725
Epoch 3/10
39/39 [==============================] - 8s 214ms/step - loss: 3.7736 - accuracy: 0.2169 - val_loss: 3.6771 - val_accuracy: 0.2279
Epoch 4/10
39/39 [==============================] - 8s 211ms/step - loss: 3.5558 - accuracy: 0.2519 - val_loss: 3.5377 - val_accuracy: 0.2499
Epoch 5/10
39/39 [==============================] - 8s 213ms/step - loss: 3.4118 - accuracy: 0.2714 - val_loss: 3.2555 - val_accuracy: 0.2810
Epoch 6/10
---------------------------------------------------------------------------
CancelledError:  [_Derived_]RecvAsync is cancelled.
     [[{{node broadcast_weights_1/assert_broadcastable/AssertGuard/else/_13/broadcast_weights_1/assert_broadcastable/AssertGuard/Assert/data_1/_74}}]] [Op:__inference_train_function_11025]

Function call stack:
train_function