You should look for the following
-
torch.cuda.empty_cache()This should free up the memory -
If the memory still does not get freed up, there is a active variable in your session that is locking up memory. Restart your session and run your code again