modified fit function to track training losses

src/dynadojo/baselines/aug_ode.py

@@ -29,7 +29,8 @@ def fit(self, x: np.ndarray, epochs=100, **kwargs):
         x = torch.tensor(x, dtype=torch.float32)
         state = x[:, 0, :]
         step = end = epochs / self._timesteps
-
+        losses = []
+
         for _ in range(epochs):
             if _ % step == 0:
                 t = torch.linspace(0.0, end, self._timesteps)
@@ -40,6 +41,10 @@ def fit(self, x: np.ndarray, epochs=100, **kwargs):
             loss = self.mse_loss(pred, x).float()
             loss.backward()
             self.opt.step()
+            losses.append(loss)
+        return {
+            "train_losses": losses
+        }
 
     def predict(self, x0, timesteps):
         x0 = torch.tensor(x0, dtype=torch.float32)