Fix SPOS state dict (#2375)

microsoft · May 8, 2020 · ef20691 · ef20691
1 parent f40242a
commit ef20691
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/nas/spos/network.py b/examples/nas/spos/network.py
@@ -147,8 +147,10 @@ def _initialize_weights(self):
 
 def load_and_parse_state_dict(filepath="./data/checkpoint-150000.pth.tar"):
     checkpoint = torch.load(filepath, map_location=torch.device("cpu"))
+    if "state_dict" in checkpoint:
+        checkpoint = checkpoint["state_dict"]
     result = dict()
-    for k, v in checkpoint["state_dict"].items():
+    for k, v in checkpoint.items():
         if k.startswith("module."):
             k = k[len("module."):]
         result[k] = v