diff --git a/mu_map/training/lib.py b/mu_map/training/lib.py index dc92337f1d3a28aea466bd39595d26bd5d09bcf6..ccb2eda79871685b398b6a0d2415554dae06b4df 100644 --- a/mu_map/training/lib.py +++ b/mu_map/training/lib.py @@ -110,7 +110,7 @@ class AbstractTraining: param.optimizer.step() return loss / len(data_loader) - def _eval_epoch(self, phase: str): + def _eval_epoch(self): torch.set_grad_enabled(False) for model in self.models: model.eval()