Let PyTorchModel define train/eval mode of LUMEModule during instanti…

…ation
slaclab · Jun 21, 2023 · dd20dc5 · dd20dc5
1 parent 861ff30
commit dd20dc5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lume_model/torch/module.py b/lume_model/torch/module.py
@@ -31,6 +31,8 @@ def __init__(
         self._feature_order = feature_order
         self._output_order = output_order
         self.register_module("base_model", self._model.model)
+        if not model.model.training:  # PyTorchModel defines train/eval mode
+            self.eval()
 
     @property
     def feature_order(self):