Keep model state

santiag0m · Mar 25, 2022 · 0ccb45a · 0ccb45a
1 parent 17e5aac
commit 0ccb45a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/lib/utils/bias.py b/lib/utils/bias.py
@@ -29,6 +29,7 @@ def reset(self):
 def compute_bias(model: nn.Module, dataloader: DataLoader) -> torch.Tensor:
     target_avg = RunningAvg()
     pred_avg = RunningAvg()
+    is_training = model.training
     model.eval()
     device = next(model.parameters()).device
     with torch.no_grad():
@@ -39,4 +40,6 @@ def compute_bias(model: nn.Module, dataloader: DataLoader) -> torch.Tensor:
             target_avg.update(targets)
             pred_avg.update(preds)
     bias = target_avg.value - pred_avg.value
+    if is_training:
+        model.train()
     return bias