Check if covariance and mapping keys match when loading state dict

curvlinops/kfac.py

@@ -1148,6 +1148,23 @@ def load_state_dict(self, state_dict: Dict[str, Any]):
         self._N_data = state_dict["num_data"]
 
         # Set Kronecker factors (if computed)
+        if self._input_covariances or self._gradient_covariances:
+            # If computed, check if the keys match the mapping keys
+            input_covariances_keys = set(self._input_covariances.keys())
+            gradient_covariances_keys = set(self._gradient_covariances.keys())
+            mapping_keys = set(self._mapping.keys())
+            if (
+                input_covariances_keys != mapping_keys
+                or gradient_covariances_keys != mapping_keys
+            ):
+                raise ValueError(
+                    "Input or gradient covariance keys in state dict do not match "
+                    "mapping keys of linear operator. "
+                    "Difference between input covariance and mapping keys: "
+                    f"{input_covariances_keys - mapping_keys}. "
+                    "Difference between gradient covariance and mapping keys: "
+                    f"{gradient_covariances_keys - mapping_keys}."
+                )
         self._input_covariances = state_dict["input_covariances"]
         self._gradient_covariances = state_dict["gradient_covariances"]
 

test/test_inverse.py

@@ -708,11 +708,11 @@ def test_KFAC_inverse_save_and_load_state_dict():
     wrong_kfac = KFACLinearOperator(model, CrossEntropyLoss(), params, [(X, y)])
     inv_kfac_wrong = KFACInverseLinearOperator(wrong_kfac)
     with raises(ValueError, match="mismatch"):
-        inv_kfac_wrong.load_state_dict(torch.load(state_dict))
+        inv_kfac_wrong.load_state_dict(torch.load("inv_kfac_state_dict.pt"))
 
     # create new inverse KFAC and load state dict
     inv_kfac_new = KFACInverseLinearOperator(kfac)
-    inv_kfac_new.load_state_dict(torch.load(state_dict))
+    inv_kfac_new.load_state_dict(torch.load("inv_kfac_state_dict.pt"))
 
     # check that the two inverse KFACs are equal
     compare_state_dicts(inv_kfac.state_dict(), inv_kfac_new.state_dict())

test/test_kfac.py

@@ -1253,7 +1253,7 @@ def test_save_and_load_state_dict():
         [(X, y)],
     )
     with raises(ValueError, match="loss"):
-        kfac_new.load_state_dict(load(state_dict))
+        kfac_new.load_state_dict(load("kfac_state_dict.pt"))
 
     # create new KFAC with different loss reduction and try to load state dict
     kfac_new = KFACLinearOperator(
@@ -1263,7 +1263,7 @@ def test_save_and_load_state_dict():
         [(X, y)],
     )
     with raises(ValueError, match="reduction"):
-        kfac_new.load_state_dict(load(state_dict))
+        kfac_new.load_state_dict(load("kfac_state_dict.pt"))
 
     # create new KFAC with different model and try to load state dict
     wrong_model = Sequential(Linear(D_in, 10), ReLU(), Linear(10, D_out))
@@ -1276,7 +1276,7 @@ def test_save_and_load_state_dict():
         loss_average=None,
     )
     with raises(RuntimeError, match="loading state_dict"):
-        kfac_new.load_state_dict(load(state_dict))
+        kfac_new.load_state_dict(load("kfac_state_dict.pt"))
 
     # create new KFAC and load state dict
     kfac_new = KFACLinearOperator(
@@ -1287,7 +1287,7 @@ def test_save_and_load_state_dict():
         loss_average=None,
         check_deterministic=False,  # turn off to avoid computing KFAC again
     )
-    kfac_new.load_state_dict(load(state_dict))
+    kfac_new.load_state_dict(load("kfac_state_dict.pt"))
 
     # check that the two KFACs are equal
     assert len(kfac.state_dict()) == len(kfac_new.state_dict())