mt5-base fine-tuning failed using single A100 GPU

[pedramyamini]

I'm trying to fine-tune mt5-base for summarization task using the script in chapter 7 of huggingface course but it fails at model.fit raising statefulpartitionedcall error using Nvidia A100 40 GB VRAM but works well for mt5-small. What is the problem? Any solution? Thank you.