From f6e651067156a08263ba22cbcbce02a0f5f4befc Mon Sep 17 00:00:00 2001 From: Joey Ballentine Date: Thu, 16 May 2024 23:30:13 -0400 Subject: [PATCH] Only call .eval if training --- libs/spandrel/spandrel/__helpers/model_descriptor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/spandrel/spandrel/__helpers/model_descriptor.py b/libs/spandrel/spandrel/__helpers/model_descriptor.py index f682a7de..c7491e19 100644 --- 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)