runtime Error:

Hi,
I am having issue running the last step of AlphaFold.ipynd on the Colab.
The error I am having is:

/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py:189: FutureWarning: jax.tree_util.tree_multimap() is deprecated. Please use jax.tree_util.tree_map() instead as a drop-in replacement.
'instead as a drop-in replacement.', FutureWarning)
---------------------------------------------------------------------------
UnfilteredStackTrace Traceback (most recent call last)
[](https://localhost:8080/#) in ()
38 processed_feature_dict = model_runner.process_features(np_example, random_seed=0)
---> 39 prediction = model_runner.predict(processed_feature_dict, random_seed=random.randrange(sys.maxsize))
40

9 frames
UnfilteredStackTrace: RuntimeError: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

RuntimeError Traceback (most recent call last)
[/opt/conda/lib/python3.7/site-packages/alphafold/model/model.py](https://localhost:8080/#) in predict(self, feat, random_seed)
165 logging.info('Running predict with shape(feat) = %s',
166 tree.map_structure(lambda x: x.shape, feat))
--> 167 result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
168
169 # This block is to ensure benchmark timings are accurate. Some blocking is

RuntimeError: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory

_**_Before this error, there was another error basically saying there is issue with running relaxing stage, so I had to disable it, but I got another error as above_**_
Thanks for your help