The torch.no_grad() guard should be used in the forward pass and include all operations, which Autograd should not track.
Could you add it to the forward and check, if this would lower the memory usage?