Running funetuning on GPU

[insundaycathy]

Thanks for releasing your great work. I was wondering if there is a way to run the finetuning and zero-shot inference code on GPU rather than TPU? What king of adjustment would I need to make?
Thanks

[rowanz]

It should be possible. The big differences might be removing the "bfloat16" casting (as that's a datatype specific to TPUs), and changing some of the boilerplate code to not be TPU specific.

[HJiashu]

Hi!
Have you been able to run the model on GPU's?

I'm also trying to run the model on GPU, but the model checkpoint is in bfloat16.
Thus I'm trying to load the model first with bfloat16's, and then cast everything back to float32's, but am encountering various problems.
I'm wondering if you have resolved this issue somehow?