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.