🔧 Fine-Tune Any HuggingFace Model like Gemma on TPUs with TorchAX

What if you could fine-tune any HuggingFace model on TPUs — using PyTorch code?


Here is what the end result looks like:



import torchax as tx
import torchax.train

# One function: forward → loss → gradients → optimizer update
step_fn = tx.train.make_train_step(model_fn, loss_fn, optimizer)

# Training loop
for batch in dataloader:
loss,...