About using with torch.no_grad()

I am new for pytorch and I had a few confusions so bear with me. I know we don’t do back propagation during the evaluation stage and for that we can use torch.no_grad(). However,I was just wondering if it’s a must to use torch.no_grad()? That is, is it okay not to use torch.no_grad() as long as we don’t do loss.backward() and optimizer.step() in evaluation step like the following:
criterion = torch.nn.CrossEntropyLoss()
for epoch in range(n_epochs):
model.train()
for x,target in data:
prediction=model(x)
optimizer.zero_grad()
loss=criterion(prediction,target)
loss.backward()
optimzer.step()

as can see from above, we do back propagation followed by step.

model.eval()
for x,target in data:
prediction = model(x)
loss=criterion(prediction,target)
VS
model.eval()
for x, target in data:
with torch.no_grad():
prediction=model(x)
loss=criterion(prediction,target)
Since we aren’t doing back propagation and updating each parameter of the model, do we necessarily have to write with torch.no_grad() and if so why is it so? thanks for your help in advance.

No, you don’t need to use the no_grad() wrapper during the evaluation phase. While it would reduce the memory usage by avoiding to store intermediate tensors needed to calculate the gradients, it’s a performance improvement and won’t change the functionality otherwise.

1 Like