If you have access to TPUs, try out the Flax training script to train even faster (this’ll also work for GPUs). With the same configuration settings, the Flax training script should be at least 70% faster than the PyTorch training script! ⚡️