Gemma colab issue

```Python
# @title Test the sampler

# Will take some time on first invocation. If it takes over a minute on a TPU
# runtime, something is probably wrong.
output = sampler.chat("Is a tomato a vegetable or a fruit?")
```

Gemma notebook doesn't work because of the known issue:

https://github.com/google-deepmind/gemma/issues/564