Make sure you run your “inference” section in a
with torch.set_grad_enabled(False):
    # your "inference" code here

Otherwise, it will be building on the autograd graph in the background and consuming more and more memory.