TPU inference
Does the provided inference code, work for TPU environment, if not?
I think it's better to include it.
If someone else has script, could you share?
Hi @Arpit-Bansal ,
Yes, the inference code works in a TPU environment. I tested it with the google/gemma-3-4b-it model, which suggests it should also be compatible with the google/gemma-3-27b-it model. Please refer to the attached gist file for details.
Thank you.
Hi @Arpit-Bansal ,
Yes, the inference code works in a TPU environment. I tested it with the
google/gemma-3-4b-itmodel, which suggests it should also be compatible with thegoogle/gemma-3-27b-itmodel. Please refer to the attached gist file for details.Thank you.
Hello, I've tried you gist file, but after executing model.device, it says cpu, instead of xla, tpu or something like that. So I doubt that the model is actually processed by CPU on the TPU machine, not TPU itself.