RuntimeError: Input tensor at index 3 has invalid shape [2, 2, 16, 128, 64] but expected [2, 4, 16, 128, 64]

Question

Runtime error while finetuning a pretrained GPT2-medium model using Huggingface library in SageMaker - ml.p3.8xlarge instance.

The finetuning_gpt2_script.py