Cuda out of memory on evaluation mode

Hello! I trained a model and I want to test it on some given values. I load the values like this:

factors = torch.from_numpy(variables)
factors = factors.cuda()
factors = factors.float()

and then I put the model in the evaluation mode and run it:

model = SimpleNet(n_variables).cuda()
model.load_state_dict(torch.load("weights.h5"))
model.eval()
model(factors)

But I am getting a CUDA error: out of memory error. If I reduce the size of “factors” it works fine. I assume I am still keeping tracks of some gradients or something, even if the model is in evaluation mode (sorry I am pretty new to pytorch). Can someone help me with this please? Thank you!

model.eval() changes the behavior of some layers, e.g. nn.Dropout won’t drop anymore and nn.BatchNorm will use the running estimates instead of the batch stats.
To save some memory you should wrap your code in with torch.no_grad(): so that the intermediate activations won’t be stored:

with torch.no_grad():
    output = model(factors)
3 Likes

Thank you so much! It works! So, why do I need to do that? What gradient is calculated during the evaluation time, if there is no loss function to be computed?

No gradients are being calculated, but the intermediate activations are stored, which are necessary to calculate the gradients. If you wrap your code in this statement, the intermediates will be cleared, which also means you cannot call .backward() on your loss anymore.

2 Likes