"RuntimeError: cannot enter context" in grpo_demo.ipynb

I attempted to run `grpo_demo.ipynb` to reproduce the training you've shown in [your blog post](https://developers.googleblog.com/en/introducing-tunix-a-jax-native-library-for-llm-post-training/).

I've provisioned a v5e-4 machine from GCP and ran the notebook.

In the code block that invokes `get_gemma_ref_model`, I get lots of error text, including:

```
ERROR:asyncio:Exception in callback Task.__step()
handle:
Traceback (most recent call last):
File "/home/ramrachum/.local/share/uv/python/cpython-3.13.9-linux-x86_64-gnu/lib/python3.13/asyncio/events.py", line 89, in _run
self._context.run(self._callback, *self._args)
~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: cannot enter context: <_contextvars.Context object at 0x7f94a16454c0> is already entered

```

I've continued to run the notebook, not sure whether this error is going to adversly affect the results. When I got to the `grpo_trainer.train` cell and ran it, I got similar exceptions:

```
ERROR:asyncio:Exception in callback Task.__step()
handle:
Traceback (most recent call last):
File "/home/ramrachum/.local/share/uv/python/cpython-3.13.9-linux-x86_64-gnu/lib/python3.13/asyncio/events.py", line 89, in _run
self._context.run(self._callback, *self._args)
~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: cannot enter context: <_contextvars.Context object at 0x7f94a16454c0> is already entered
```

Here is an uploaded copy of the notebook showing the error: https://colab.research.google.com/drive/1cMd3jKUiyUUGBvMELtaPthdqIwNImdNy

Note that I did not run this on Colab, I ran it on Jupyterlab on a GCP machine, because Colab does not provide a v5e-4 machine as far as I know.

**Environment**

- **OS:** Ubuntu 22.04.3 LTS
- **Project Version:** `git+https://github.com/google/tunix@f0d5d1e63e24f42647fd4e6122641d689f8bfd0e`

**Checklist**

- [X] I have searched the existing issues for a similar bug report.
- [X] I have provided all the required information in the "Environment" section.
- [X] I have provided a minimal, reproducible example.

**Would you like to help us fix it?**

Yes.