I made it using colab trainer in a complicated way that I can't remember.