boolean mode in module.train

szagoruyko · soumith · commit b5f75921404b · 2017-03-02T09:18:05.000-05:00
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
@@ -357,25 +357,22 @@ def modules(self, memo=None):
                 for m in module.modules(memo):
                     yield m
 
-    def train(self):
+    def train(self, mode=True):
         """Sets the module in training mode.
 
         This has any effect only on modules such as Dropout or BatchNorm.
         """
-        self.training = True
+        self.training = mode
         for module in self.children():
-            module.train()
+            module.train(mode)
         return self
 
     def eval(self):
         """Sets the module in evaluation mode.
 
         This has any effect only on modules such as Dropout or BatchNorm.
         """
-        self.training = False
-        for module in self.children():
-            module.eval()
-        return self
+        return self.train(False)
 
     def zero_grad(self):
         """Sets gradients of all model parameters to zero."""
