To reproduce the problem run https://github.com/NSAPH-Projects/transformxclimate/blob/main/gpt_training.ipynb