Getting jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED error while performing inference on gemma-2b on TPU

I ran the whole colab script on my TPU server - https://colab.research.google.com/github/google-deepmind/gemma/blob/main/colabs/sampling_tutorial.ipynb#scrollTo=tqbJ1SUcESaN

The script ran properly the first time, but sometime later, when I again ran the script I got the following error

`jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 15.65G. That was not possible. There are 11.08G free.; (0x0x0_HBM0)`

Can anyone please tell what is the exact cause of the error and how to fix this?