added a warning when a PyTorch model is in train mode (#78)

* added a warning when a PyTorch model is in train mode * added missing import
bethgelab · Oct 6, 2017 · 76407ec · 76407ec
1 parent fcbbc6f
commit 76407ec
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/foolbox/models/pytorch.py b/foolbox/models/pytorch.py
@@ -1,4 +1,5 @@
 import numpy as np
+import warnings
 
 from .base import DifferentiableModel
 
@@ -43,6 +44,12 @@ def __init__(
         self._model = model
         self.cuda = cuda
 
+        if model.training:
+            warnings.warn(
+                'The PyTorch model is in training mode and therefore might'
+                ' not be deterministic. Call the eval() method to set it in'
+                ' evaluation mode if this is not intended.')
+
     def batch_predictions(self, images):
         # lazy import
         import torch