Gemma 3 1b model loading asyncio error on kaggle notebook

**Expected Behavior**
The model loading without any asyncio error.
**Actual Behavior**
The model is loaded properly. But the notebook cell execution is stopped after the model load cell. Need to manually run after the model loading cells.

**Steps to Reproduce the Problem**

[tunix-error.ipynb](https://github.com/user-attachments/files/24148738/tunix-error.ipynb)

**Environment**

- **OS:** Kaggle TPU Notebook
- **Project Version:**

| Package Name | Version |
| :--- | :--- |
| jax[tpu] | 0.8.1 |
| flax | 0.12.1 |
| optax | 0.2.6 |
| qwix | 0.1.5 |
| orbax-checkpoint | 0.11.31 |
| grain | 0.2.15 |
| google-metrax | 0.2.4 |
| google-tunix[prod] | 0.1.5 |

**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?**