reverted change

Novartis · Jul 29, 2024 · 3e3051a · 3e3051a
1 parent 42a95a5
commit 3e3051a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/torchsurv/loss/momentum.py b/src/torchsurv/loss/momentum.py
@@ -183,8 +183,8 @@ def infer(self, inputs: torch.Tensor) -> torch.Tensor:
                     [ 0.9771, -0.8513]])
 
         """
-        with self.target.eval():  # notify all your layers that you are in eval mode
-            return self.target(inputs)
+        self.target.eval()  # notify all your layers that you are in eval mode
+        return self.target(inputs)
 
     def _bank_loss(self) -> torch.Tensor:
         """computer the  negative loss likelyhood from memory bank"""