CUDA out of memory even when memory is free?

Hello guys,
I have been running a github repository that started showing an abrupt behaviour. Some times the code worked without any error and other times it threw the below error.
RuntimeError: CUDA out of memory. Tried to allocate 2.26 GiB (GPU 0; 6.00 GiB total capacity; 53.78 MiB already allocated; 4.41 GiB free; 8.22 MiB cached)
I am running the code from the following repository:

Repo:

Command I run:
python main.py --ae --getz