Fix loading encoder weights trained with BYOL

microsoft · May 3, 2022 · a39bad6 · a39bad6
1 parent cf2a0cc
commit a39bad6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchgeo/trainers/utils.py b/torchgeo/trainers/utils.py
@@ -42,8 +42,8 @@ def extract_encoder(path: str) -> Tuple[str, "OrderedDict[str, Tensor]"]:
         state_dict = OrderedDict(
             {k.replace("model.", ""): v for k, v in state_dict.items()}
         )
-    elif "encoder" in checkpoint["hyper_parameters"]:
-        name = checkpoint["hyper_parameters"]["encoder"]
+    elif "encoder_name" in checkpoint["hyper_parameters"]:
+        name = checkpoint["hyper_parameters"]["encoder_name"]
         state_dict = checkpoint["state_dict"]
         state_dict = OrderedDict(
             {k: v for k, v in state_dict.items() if "model.encoder.model" in k}