Fix tests

curvlinops/kfac.py

@@ -1089,7 +1089,7 @@ def state_dict(self) -> Dict[str, Any]:
             "loss_reduction": self._loss_func.reduction,
             # Attributes
             "progressbar": self._progressbar,
-            "shape": self._shape,
+            "shape": self.shape,
             "seed": self._seed,
             "fisher_type": self._fisher_type,
             "mc_samples": self._mc_samples,
@@ -1137,7 +1137,7 @@ def load_state_dict(self, state_dict: Dict[str, Any]):
 
         # Set attributes
         self._progressbar = state_dict["progressbar"]
-        self._shape = state_dict["shape"]
+        self.shape = state_dict["shape"]
         self._seed = state_dict["seed"]
         self._fisher_type = state_dict["fisher_type"]
         self._mc_samples = state_dict["mc_samples"]
@@ -1182,6 +1182,9 @@ def from_state_dict(
 
         Returns:
             Linear operator of KFAC approximation.
+
+        Raises:
+            RuntimeError: If the check for deterministic behavior fails.
         """
         loss_func = {
             "MSELoss": MSELoss,

test/test_inverse.py

@@ -683,10 +683,10 @@ def test_KFAC_inverse_save_and_load_state_dict():
     # save state dict
     state_dict = inv_kfac.state_dict()
 
+    # create new inverse KFAC with different linop input and try to load state dict
+    wrong_kfac = KFACLinearOperator(model, CrossEntropyLoss(), params, [(X, y)])
+    inv_kfac_wrong = KFACInverseLinearOperator(wrong_kfac)
     with raises(ValueError, match="mismatch"):
-        # create new inverse KFAC with different linop input and try to load state dict
-        wrong_kfac = KFACLinearOperator(model, CrossEntropyLoss(), params, [(X, y)])
-        inv_kfac_wrong = KFACInverseLinearOperator(wrong_kfac)
         inv_kfac_wrong.load_state_dict(state_dict)
 
     # create new inverse KFAC and load state dict

test/test_kfac.py

@@ -1243,37 +1243,37 @@ def test_save_and_load_state_dict():
     # save state dict
     state_dict = kfac.state_dict()
 
+    # create new KFAC with different loss function and try to load state dict
+    kfac_new = KFACLinearOperator(
+        model,
+        CrossEntropyLoss(),
+        params,
+        [(X, y)],
+    )
     with raises(ValueError, match="loss"):
-        # create new KFAC with different loss function and try to load state dict
-        kfac_new = KFACLinearOperator(
-            model,
-            CrossEntropyLoss(),
-            params,
-            [(X, y)],
-        )
         kfac_new.load_state_dict(state_dict)
 
+    # create new KFAC with different loss reduction and try to load state dict
+    kfac_new = KFACLinearOperator(
+        model,
+        MSELoss(),
+        params,
+        [(X, y)],
+    )
     with raises(ValueError, match="reduction"):
-        # create new KFAC with different loss reduction and try to load state dict
-        kfac_new = KFACLinearOperator(
-            model,
-            MSELoss(),
-            params,
-            [(X, y)],
-        )
         kfac_new.load_state_dict(state_dict)
 
+    # create new KFAC with different model and try to load state dict
+    wrong_model = Sequential(Linear(D_in, 10), ReLU(), Linear(10, D_out))
+    wrong_params = list(wrong_model.parameters())
+    kfac_new = KFACLinearOperator(
+        wrong_model,
+        MSELoss(reduction="sum"),
+        wrong_params,
+        [(X, y)],
+        loss_average=None,
+    )
     with raises(RuntimeError, match="loading state_dict"):
-        # create new KFAC with different model and try to load state dict
-        wrong_model = Sequential(Linear(D_in, 10), ReLU(), Linear(10, D_out))
-        wrong_params = list(wrong_model.parameters())
-        kfac_new = KFACLinearOperator(
-            wrong_model,
-            MSELoss(reduction="sum"),
-            wrong_params,
-            [(X, y)],
-            loss_average=None,
-        )
         kfac_new.load_state_dict(state_dict)
 
     # create new KFAC and load state dict