Fix backprop of pytorch wrapper
@@ -31,9 +31,11 @@ class PyTT_Wrapper(PyTorchWrapper):
ids = xp2torch(self.ops.asarray(ids))
is_training = self._model.training
if drop is None:
+ with torch.autograd.set_grad_enabled(False):
self._model.eval()
y_var = self._model(ids)
else:
+ with torch.autograd.set_grad_enabled(True):
self._model.train()
y_var = self._model(ids)
self._model.training = is_training
