Modified XL Lora trainer from https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL.ipynb#scrollTo=OglZzI_ujZq-