with no_grad in evaluation step.
@@ -1205,6 +1205,9 @@ class TorchAgent(Agent):
if is_training:
output = self.train_step(batch)
else:
+ with torch.no_grad():
+ # save memory and compute by disabling autograd.
+ # use `with torch.enable_grad()` to gain back graidients.
output = self.eval_step(batch)
if output is None:
