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