Run Vision Transformer code (based on jax) on TPU

[lucasliunju]

Hi all,

I'm trying to run the code (https://github.com/google-research/vision_transformer) on tpu and I find it doesn't work and occurs "segmentation fault".

To run the code on tpu, I add this code on the main function:

from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://[My TPU ip address]:8470"

I'm not sure whether the code of using tpu is right.

Thank you very much!