interactive_eager_few_shot_od_training_colab.ipynb error: Failed to get convolution algorithm.

Prerequisites

• I am using the latest TensorFlow Model Garden release and TensorFlow 2.
• I am reporting the issue to the correct repository. (Model Garden official or research directory)
• I checked to make sure that this issue has not been filed already.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/blob/master/research/object_detection/colab_tutorials/eager_few_shot_od_training_tflite.ipynb

2. Describe the bug

Running this notebook on Colab with GPU enabled, on calling detection_model.predict the following error is raised:

UnknownError                              Traceback (most recent call last)
<ipython-input-23-5597398cc2f6> in <module>()
     39 # Run model through a dummy image so that variables are created
     40 image, shapes = detection_model.preprocess(tf.zeros([1, 640, 640, 3]))
---> 41 prediction_dict = detection_model.predict(image, shapes)
     42 _ = detection_model.postprocess(prediction_dict, shapes)
     43 print('Weights restored!')

15 frames
/usr/local/lib/python3.7/dist-packages/six.py in raise_from(value, from_value)

UnknownError: Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above. [Op:Conv2D]

3. Steps to reproduce

Run the notebook

4. Expected behavior

No error

5. Additional context

None

6. System information

• Colab
• The notebook uses a pinned version of tf