How to enable BF16 on tpus?

With notebook_launcher(main, use_fp16=True) my data still fp32

YEs that would be expected, this does not control bfloat16. Support for bfloat16 on TPUs is not in Accelerate yet.

Any way I can do so? Torch xla docs say to use env var, but it doesn’t work too.