Make save/load cleaner

Author: CeliaBenquet

cebra/solver/base.py

@@ -296,7 +296,7 @@ def state_dict(self) -> dict:
             the model was trained with.
         """
 
-        return {
+        state_dict = {
             "model": self.model.state_dict(),
             "optimizer": self.optimizer.state_dict(),
             "loss": torch.tensor(self.history),
@@ -306,6 +306,13 @@ def state_dict(self) -> dict:
             "log": self.log,
         }
 
+        if hasattr(self, "n_features"):
+            state_dict["n_features"] = self.n_features
+        if hasattr(self, "num_sessions"):
+            state_dict["num_sessions"] = self.num_sessions
+
+        return state_dict
+
     def load_state_dict(self, state_dict: dict, strict: bool = True):
         """Update the solver state with the given state_dict.
 
@@ -343,6 +350,12 @@ def _get(key):
         if _contains("log"):
             self.log = _get("log")
 
+        # Not defined if the model was saved before being fitted.
+        if "n_features" in state_dict:
+            self.n_features = _get("n_features")
+        if "num_sessions" in state_dict:
+            self.num_sessions = _get("num_sessions")
+
     @property
     def num_parameters(self) -> int:
         """Total number of parameters in the encoder and criterion."""
@@ -633,11 +646,6 @@ def load(self, logdir, filename="checkpoint.pth"):
         checkpoint = torch.load(savepath, map_location=self.device)
         self.load_state_dict(checkpoint, strict=True)
 
-        n_features = self.n_features
-        self.n_features = ([
-            session_n_features for session_n_features in n_features
-        ] if isinstance(n_features, list) else n_features)
-
     def save(self, logdir, filename="checkpoint.pth"):
         """Save the model and optimizer params.