GPU memory use on train.py

Thanks for you fascinating job! 
I'm trying to run the giving train code on 4060ti 16GB GPU. But it ends up with "CUDA out of memory". I wonder how large will the memory be enough to run the train.py code? I noticed that in #13 it was mentioned that it took a little GPU memory. or is there some magic that can optimize the GPU memory use?
Thank you!