The problem is, (talkin by experience), is that we save cuda weights. However, when we load the state dict, we open those weights on the gpu and memory is never freed.
If i’m not wrong, opening weights this way should work (at least talking about gpu vram, never checked cpu ram):
state_dict = torch.load(directory, map_location=lambda storage, loc: storage)
so then you can load that state dict