From 3e3051ac1ef9bcf4ebade82856ad18ea9ebcb13a Mon Sep 17 00:00:00 2001
From: corolth1 <thibaud.coroller@novartis.com>
Date: Mon, 29 Jul 2024 14:08:28 -0400
Subject: [PATCH] reverted change

---
 src/torchsurv/loss/momentum.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/torchsurv/loss/momentum.py b/src/torchsurv/loss/momentum.py
index 545cbf0..8cab479 100644
--- 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"""