Jiewen Tan

> > So, with current nightly I do not see errors anymore if running without compilation. With compilation enabled, I get the same error I see in the issue description....

@JackCaoG Basing on the latest reply from @tengomucho, it seems like a dynamo issue. Can you take a look as well?

@PawKanarek libtpu is a pip pkg, you can grep it from pip list. The latest version is: ``` pip list | grep libtpu libtpu-nightly 0.1.dev20240213 ``` If yours is older...

@PawKanarek Could be a hardware issue then... Can you try recreate a new TPU vm?

@JackCaoG We use `openxla_eval` by default in most of examples. lol We can re-benchmark it to see if the performance gaps are gone.

@PawKanarek For Gemma, have you set the following env: PJRT_DEVICE=TPU XLA_USE_SPMD=1 ?

> It seems that setting `export PJRT_DEVICE=TPU` and `export XLA_USE_SPMD=1` resolved the issue. I was certain I had exported the variables... The training now works though it occasionally crashes during...

@windmaple @PawKanarek Are we good to close this issue?

Could that buffer be donated?

Do you mind adding a test case?