train: track learn rate

MPoL-dev · Nov 29, 2023 · 16e9499 · 16e9499
1 parent bacdeec
commit 16e9499
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/mpol/training.py b/src/mpol/training.py
@@ -223,7 +223,7 @@ def train(self, model, dataset):
 
         count = 0
         losses = []
-        self._train_state = {}
+        learn_rates = []
 
         # guess initial strengths for regularizers in `self._regularizers`
         # that have 'guess':True
@@ -275,6 +275,7 @@ def train(self, model, dataset):
 
             if self._scheduler is not None:
                 self._scheduler.step(loss)
+                learn_rates.append(self._optimizer.param_groups[0]['lr'])
 
             # generate optional fit diagnostics
             if self._train_diag_step is not None and (count % self._train_diag_step == 0 or count == self._epochs or self.loss_convergence(np.array(losses))):