Hi @Sahajtomar thanks for reaching on the forum!

Are you using your own for loop or the HF Trainer?
Did you look at the T5 Distributed training example released few weeks ago? Distributed Training: Train BART/T5 for Summarization using 🤗 Transformers and Amazon SageMaker 130