[FEATURE] Inference on TPU is not working (bits_and_tpu branch)

[eliahuhorwitz]

Hey, I've been using the bits_and_tpu branch to train on TPU VMs, however when I try to run the inference script it needs a GPU, could the inference code be easily adjusted to work on TPU?

Thanks, Eliahu

[rwightman]

@eliahuhorwitz yes, it's not a huge task but it's not anywhere on my priority list right now, I still have a lot of other things to work on for training, etc. Open to PR as all the needed changes can be references from the validation script.

[eliahuhorwitz]

Thanks