Cuda out of memory when customizing model

Link to code:

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)