No TPU found, falling back to CPU

[0x7o]

I'm trying to convert the model with the "to_hf_weights.py" script. I need to load the model on the TPU. For this purpose, I use Google Colab.

The script shows a warning. WARNING:absl:No GPU/TPU found, falling back to CPU.

At the same time the script shows that the TPU is working.

import jax jax.devices()

[TpuDevice(id=0, host_id=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, host_id=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, host_id=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, host_id=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, host_id=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, host_id=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, host_id=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, host_id=0, coords=(1,1,0), core_on_chip=1)]

[kingoflolz]

to_hf_weights has not been tested on colab, and is designed for use on TPU-VMs (and hence does not initialize the TPUs properly for colab)