Minor fix in checkpoint loading

cvg · Mar 28, 2023 · 3d7bcd6 · 3d7bcd6
1 parent 0df135f
commit 3d7bcd6
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/sold2/misc/train_utils.py b/sold2/misc/train_utils.py
@@ -53,11 +53,13 @@ def remove_old_checkpoints(checkpoint_root, max_ckpt=15):
 
 
 def adapt_checkpoint(state_dict):
+    new_state_dict = {}
     for k, v in state_dict.items():
         if k.startswith('module.'):
-            state_dict[k[7:]] = v
-            del state_dict[k]
-    return state_dict
+            new_state_dict[k[7:]] = v
+        else:
+            new_state_dict[k] = v
+    return new_state_dict
 
 
 ################