Cuda out of memory when customizing model

Link to code: https://colab.research.google.com/drive/14nqc0qB1VdKKbiPidg9KSsnO7SBEE-yV?usp=sharing

It always give me a cuda out of memory error at 4th epoch or in the 5th epoch.

Error: CUDA out of memory. Tried to allocate 224.00 MiB (GPU 0; 11.17 GiB total capacity; 10.30 GiB already allocated; 222.81 MiB free; 10.53 GiB reserved in total by PyTorch)