Fix save/load

Author: CeliaBenquet

cebra/solver/base.py

@@ -668,13 +668,12 @@ def load(self, logdir, filename="checkpoint.pth"):
         checkpoint = torch.load(savepath, map_location=self.device)
         self.load_state_dict(checkpoint, strict=True)
 
-        if hasattr(self.model, "n_features"):
-            n_features = self.model.n_features
-            self.n_features = ([
-                session_n_features for session_n_features in n_features
-            ] if isinstance(n_features, list) else n_features)
+        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_last.pth"):
+    def save(self, logdir, filename="checkpoint.pth"):
         """Save the model and optimizer params.
 
         Args:

tests/test_solver.py

@@ -20,6 +20,7 @@
 # limitations under the License.
 #
 import copy
+import itertools
 import tempfile
 
 import numpy as np
@@ -189,6 +190,12 @@ def test_single_session(data_name, loader_initfunc, model_architecture,
     for param in solver.parameters():
         assert isinstance(param, torch.Tensor)
 
+    fitted_solver = copy.deepcopy(solver)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        solver.save(temp_dir)
+        solver.load(temp_dir)
+    _assert_equal(fitted_solver, solver)
+
 
     embedding = solver.transform(X)
     assert isinstance(embedding, torch.Tensor)
@@ -380,6 +387,12 @@ def test_multi_session(data_name, loader_initfunc, model_architecture,
         for param in solver.parameters():
             assert isinstance(param, torch.Tensor)
 
+    fitted_solver = copy.deepcopy(solver)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        solver.save(temp_dir)
+        solver.load(temp_dir)
+    _assert_equal(fitted_solver, solver)
+
 
 @pytest.mark.parametrize(
     "inputs, add_padding, offset, start_batch_idx, end_batch_idx, expected_output",