Skip to content

GitLab

Unverified Commit ef206913 authored by Yuge Zhang's avatar Yuge Zhang Committed by GitHub
Browse files

Fix SPOS state dict (#2375)

parent f40242a7
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
examples/nas/spos/network.py
View file @ ef206913
...@@ -147,8 +147,10 @@ class ShuffleNetV2OneShot(nn.Module): ...@@ -147,8 +147,10 @@ class ShuffleNetV2OneShot(nn.Module):
def load_and_parse_state_dict(filepath="./data/checkpoint-150000.pth.tar"): def load_and_parse_state_dict(filepath="./data/checkpoint-150000.pth.tar"):
checkpoint = torch.load(filepath, map_location=torch.device("cpu")) checkpoint = torch.load(filepath, map_location=torch.device("cpu"))
if "state_dict" in checkpoint:
checkpoint = checkpoint["state_dict"]
result = dict() result = dict()
for k, v in checkpoint["state_dict"].items(): for k, v in checkpoint.items():
if k.startswith("module."): if k.startswith("module."):
k = k[len("module."):] k = k[len("module."):]
result[k] = v result[k] = v
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment