Only call .eval if training

chaiNNer-org · May 17, 2024 · f6e6510 · f6e6510
1 parent 65ebd7e
commit f6e6510
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/libs/spandrel/spandrel/__helpers/model_descriptor.py b/libs/spandrel/spandrel/__helpers/model_descriptor.py
@@ -465,7 +465,8 @@ def __call__(self, image: Tensor) -> Tensor:
         did_pad, image = pad_tensor(image, self.size_requirements)
 
         # Optimize for inference
-        self.model.eval()
+        if self.model.training:
+            self.model.eval()
 
         # call model
         output = self._call_fn(self.model, image)