Hello,
I don’t know if this can solve your issue. But I was having a similar one (i am looping a model over different folds) and at the second loop I obtained this famous memory error.
At the end of the loop I free the memory in the following way:

for element in dir():
    if element[0:2] != "__":

        del globals()[element]

import torch

torch.cuda.empty_cache()


I found this solution on https://www.kite.com/python/answers/how-to-delete-variables-and-functions-from-memory-in-python 2
I hope it solves your problem