JIT compile always crashes the kernel and restarts on google colab TPU.

[anurag12-webster]

When i execute below code the kernel crashes everytime and restarts.

code:

from whisper_jax import FlaxWhisperPipline

# instantiate pipeline
pipeline = FlaxWhisperPipline("openai/whisper-large-v2")

# JIT compile the forward call - slow, but we only do once
text = pipeline("elon8m.mp3")

# used cached function thereafter - super fast!!
text = pipeline("elon8m.mp3")

console output

/usr/local/lib/python3.10/dist-packages/transformers/utils/generic.py:311: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  torch.utils._pytree._register_pytree_node(

below is the image for reference

[quantumsleeper]

Same error on Kaggle!