Merge pull request #1 from oneTaken/master

keep compatibility when no parallel
he-y · Jun 14, 2019 · 858ebac · 858ebac
2 parents 0290ab0 + cbd2399
commit 858ebac
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/functions/infer_pruned.py b/functions/infer_pruned.py
@@ -210,10 +210,10 @@ def accuracy(output, target, topk=(1,)):
 def remove_module_dict(state_dict):
     new_state_dict = OrderedDict()
     for k, v in state_dict.items():
-        name = k[7:]  # remove `module.`
+        name = name[7:] if name.startswith("module.") else name  # remove `module.`
         new_state_dict[name] = v
     return new_state_dict
 
 
 if __name__ == '__main__':
-    main()
+    main()