Fixed loss (we were saving model with best 'loss' rather than best 't…

…otal_loss')
arnauqb · Sep 17, 2024 · 4800a31 · 4800a31
1 parent b31d0b0
commit 4800a31
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/blackbirds/infer/vi.py b/blackbirds/infer/vi.py
@@ -507,8 +507,8 @@ def run(
                         "Loss/regularisation", regularisation_loss, epoch
                     )
                 torch.save(self.best_estimator_state_dict, "last_estimator.pt")
-                if loss < self.best_loss:
-                    self.best_loss = loss
+                if total_loss < self.best_loss:
+                    self.best_loss = total_loss
                     self.best_estimator_state_dict = deepcopy(
                         self.posterior_estimator.state_dict()
                     )