GPU Memory in Eval vs Training

Your assumption is correct. During evaluation you don’t need to calculate the gradients, which would take memory, and also don’t need to store intermediate activations, which would be needed to calculate the gradients.

To lower the memory usage and not store these intermediates, you should wrap your evaluation code into a with torch.no_grad() block as seen here:

model = MyModel().to('cuda')
with torch.no_grad():
    output = model(data)
1 Like