Set no_grad for prediction in wrapper
@@ -48,6 +48,7 @@ class PyTT_Wrapper(PyTorchWrapper):
model_kwargs = self.get_model_kwargs(ids)
if drop is None:
self._model.eval()
+ with torch.no_grad():
y_var = self._model(ids, **model_kwargs)
else:
self._model.train()
