Try being explicit

microsoft · Jan 31, 2024 · b5a9007 · b5a9007
1 parent 35b914c
commit b5a9007
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torchgeo/trainers/utils.py b/torchgeo/trainers/utils.py
@@ -117,7 +117,10 @@ def load_state_dict(
             state_dict[output_module_key + ".bias"],
         )
 
-    return model.load_state_dict(state_dict, strict=False)
+    missing_keys: list[str]
+    unexpected_keys: list[str]
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    return missing_keys, unexpected_keys
 
 
 def reinit_initial_conv_layer(