We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
BUG in spos evolution part
I trained the spos from scratch with supernet.py, but failed to load the checkpoints. It turned out there was a mismatch between load and save part. 

In  the save part  https://github.com/microsoft/nni/blob/69dfbf5e976cf75e89949af4764928dde631a3f6/src/sdk/pynni/nni/nas/pytorch/callbacks.py#L136,  no dict was used. But in the load part,  https://github.com/microsoft/nni/blob/69dfbf5e976cf75e89949af4764928dde631a3f6/examples/nas/spos/network.py#L151,  key 'state_dict' was used. It failed to load the checkpoint. After removing 'state_dict', checkpoint could be loaded and problem was solved.



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/nas/spos/network.py`
Content:
```
1 # Copyright (c) Microsoft Corporation.
2 # Licensed under the MIT license.
3 
4 import os
5 import pickle
6 import re
7 
8 import torch
9 import torch.nn as nn
10 from nni.nas.pytorch import mutables
11 
12 from blocks import ShuffleNetBlock, ShuffleXceptionBlock
13 
14 
15 class ShuffleNetV2OneShot(nn.Module):
16     block_keys = [
17         'shufflenet_3x3',
18         'shufflenet_5x5',
19         'shufflenet_7x7',
20         'xception_3x3',
21     ]
22 
23     def __init__(self, input_size=224, first_conv_channels=16, last_conv_channels=1024, n_classes=1000,
24                  op_flops_path="./data/op_flops_dict.pkl"):
25         super().__init__()
26 
27         assert input_size % 32 == 0
28         with open(os.path.join(os.path.dirname(__file__), op_flops_path), "rb") as fp:
29             self._op_flops_dict = pickle.load(fp)
30 
31         self.stage_blocks = [4, 4, 8, 4]
32         self.stage_channels = [64, 160, 320, 640]
33         self._parsed_flops = dict()
34         self._input_size = input_size
35         self._feature_map_size = input_size
36         self._first_conv_channels = first_conv_channels
37         self._last_conv_channels = last_conv_channels
38         self._n_classes = n_classes
39 
40         # building first layer
41         self.first_conv = nn.Sequential(
42             nn.Conv2d(3, first_conv_channels, 3, 2, 1, bias=False),
43             nn.BatchNorm2d(first_conv_channels, affine=False),
44             nn.ReLU(inplace=True),
45         )
46         self._feature_map_size //= 2
47 
48         p_channels = first_conv_channels
49         features = []
50         for num_blocks, channels in zip(self.stage_blocks, self.stage_channels):
51             features.extend(self._make_blocks(num_blocks, p_channels, channels))
52             p_channels = channels
53         self.features = nn.Sequential(*features)
54 
55         self.conv_last = nn.Sequential(
56             nn.Conv2d(p_channels, last_conv_channels, 1, 1, 0, bias=False),
57             nn.BatchNorm2d(last_conv_channels, affine=False),
58             nn.ReLU(inplace=True),
59         )
60         self.globalpool = nn.AvgPool2d(self._feature_map_size)
61         self.dropout = nn.Dropout(0.1)
62         self.classifier = nn.Sequential(
63             nn.Linear(last_conv_channels, n_classes, bias=False),
64         )
65 
66         self._initialize_weights()
67 
68     def _make_blocks(self, blocks, in_channels, channels):
69         result = []
70         for i in range(blocks):
71             stride = 2 if i == 0 else 1
72             inp = in_channels if i == 0 else channels
73             oup = channels
74 
75             base_mid_channels = channels // 2
76             mid_channels = int(base_mid_channels)  # prepare for scale
77             choice_block = mutables.LayerChoice([
78                 ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=3, stride=stride),
79                 ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=5, stride=stride),
80                 ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=7, stride=stride),
81                 ShuffleXceptionBlock(inp, oup, mid_channels=mid_channels, stride=stride)
82             ])
83             result.append(choice_block)
84 
85             # find the corresponding flops
86             flop_key = (inp, oup, mid_channels, self._feature_map_size, self._feature_map_size, stride)
87             self._parsed_flops[choice_block.key] = [
88                 self._op_flops_dict["{}_stride_{}".format(k, stride)][flop_key] for k in self.block_keys
89             ]
90             if stride == 2:
91                 self._feature_map_size //= 2
92         return result
93 
94     def forward(self, x):
95         bs = x.size(0)
96         x = self.first_conv(x)
97         x = self.features(x)
98         x = self.conv_last(x)
99         x = self.globalpool(x)
100 
101         x = self.dropout(x)
102         x = x.contiguous().view(bs, -1)
103         x = self.classifier(x)
104         return x
105 
106     def get_candidate_flops(self, candidate):
107         conv1_flops = self._op_flops_dict["conv1"][(3, self._first_conv_channels,
108                                                     self._input_size, self._input_size, 2)]
109         # Should use `last_conv_channels` here, but megvii insists that it's `n_classes`. Keeping it.
110         # https://github.com/megvii-model/SinglePathOneShot/blob/36eed6cf083497ffa9cfe7b8da25bb0b6ba5a452/src/Supernet/flops.py#L313
111         rest_flops = self._op_flops_dict["rest_operation"][(self.stage_channels[-1], self._n_classes,
112                                                             self._feature_map_size, self._feature_map_size, 1)]
113         total_flops = conv1_flops + rest_flops
114         for k, m in candidate.items():
115             parsed_flops_dict = self._parsed_flops[k]
116             if isinstance(m, dict):  # to be compatible with classical nas format
117                 total_flops += parsed_flops_dict[m["_idx"]]
118             else:
119                 total_flops += parsed_flops_dict[torch.max(m, 0)[1]]
120         return total_flops
121 
122     def _initialize_weights(self):
123         for name, m in self.named_modules():
124             if isinstance(m, nn.Conv2d):
125                 if 'first' in name:
126                     nn.init.normal_(m.weight, 0, 0.01)
127                 else:
128                     nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
129                 if m.bias is not None:
130                     nn.init.constant_(m.bias, 0)
131             elif isinstance(m, nn.BatchNorm2d):
132                 if m.weight is not None:
133                     nn.init.constant_(m.weight, 1)
134                 if m.bias is not None:
135                     nn.init.constant_(m.bias, 0.0001)
136                 nn.init.constant_(m.running_mean, 0)
137             elif isinstance(m, nn.BatchNorm1d):
138                 nn.init.constant_(m.weight, 1)
139                 if m.bias is not None:
140                     nn.init.constant_(m.bias, 0.0001)
141                 nn.init.constant_(m.running_mean, 0)
142             elif isinstance(m, nn.Linear):
143                 nn.init.normal_(m.weight, 0, 0.01)
144                 if m.bias is not None:
145                     nn.init.constant_(m.bias, 0)
146 
147 
148 def load_and_parse_state_dict(filepath="./data/checkpoint-150000.pth.tar"):
149     checkpoint = torch.load(filepath, map_location=torch.device("cpu"))
150     result = dict()
151     for k, v in checkpoint["state_dict"].items():
152         if k.startswith("module."):
153             k = k[len("module."):]
154         result[k] = v
155     return result
156
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/nas/spos/network.py b/examples/nas/spos/network.py
--- a/examples/nas/spos/network.py
+++ b/examples/nas/spos/network.py
@@ -147,8 +147,10 @@
 
 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

{"golden_diff": "diff --git a/examples/nas/spos/network.py b/examples/nas/spos/network.py\n--- a/examples/nas/spos/network.py\n+++ b/examples/nas/spos/network.py\n@@ -147,8 +147,10 @@\n \n def load_and_parse_state_dict(filepath=\"./data/checkpoint-150000.pth.tar\"):\n     checkpoint = torch.load(filepath, map_location=torch.device(\"cpu\"))\n+    if \"state_dict\" in checkpoint:\n+        checkpoint = checkpoint[\"state_dict\"]\n     result = dict()\n-    for k, v in checkpoint[\"state_dict\"].items():\n+    for k, v in checkpoint.items():\n         if k.startswith(\"module.\"):\n             k = k[len(\"module.\"):]\n         result[k] = v\n", "issue": "BUG in spos evolution part\nI trained the spos from scratch with supernet.py, but failed to load the checkpoints. It turned out there was a mismatch between load and save part. \r\n\r\nIn  the save part  https://github.com/microsoft/nni/blob/69dfbf5e976cf75e89949af4764928dde631a3f6/src/sdk/pynni/nni/nas/pytorch/callbacks.py#L136,  no dict was used. But in the load part,  https://github.com/microsoft/nni/blob/69dfbf5e976cf75e89949af4764928dde631a3f6/examples/nas/spos/network.py#L151,  key 'state_dict' was used. It failed to load the checkpoint. After removing 'state_dict', checkpoint could be loaded and problem was solved.\r\n\r\n\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# Licensed under the MIT license.\n\nimport os\nimport pickle\nimport re\n\nimport torch\nimport torch.nn as nn\nfrom nni.nas.pytorch import mutables\n\nfrom blocks import ShuffleNetBlock, ShuffleXceptionBlock\n\n\nclass ShuffleNetV2OneShot(nn.Module):\n    block_keys = [\n        'shufflenet_3x3',\n        'shufflenet_5x5',\n        'shufflenet_7x7',\n        'xception_3x3',\n    ]\n\n    def __init__(self, input_size=224, first_conv_channels=16, last_conv_channels=1024, n_classes=1000,\n                 op_flops_path=\"./data/op_flops_dict.pkl\"):\n        super().__init__()\n\n        assert input_size % 32 == 0\n        with open(os.path.join(os.path.dirname(__file__), op_flops_path), \"rb\") as fp:\n            self._op_flops_dict = pickle.load(fp)\n\n        self.stage_blocks = [4, 4, 8, 4]\n        self.stage_channels = [64, 160, 320, 640]\n        self._parsed_flops = dict()\n        self._input_size = input_size\n        self._feature_map_size = input_size\n        self._first_conv_channels = first_conv_channels\n        self._last_conv_channels = last_conv_channels\n        self._n_classes = n_classes\n\n        # building first layer\n        self.first_conv = nn.Sequential(\n            nn.Conv2d(3, first_conv_channels, 3, 2, 1, bias=False),\n            nn.BatchNorm2d(first_conv_channels, affine=False),\n            nn.ReLU(inplace=True),\n        )\n        self._feature_map_size //= 2\n\n        p_channels = first_conv_channels\n        features = []\n        for num_blocks, channels in zip(self.stage_blocks, self.stage_channels):\n            features.extend(self._make_blocks(num_blocks, p_channels, channels))\n            p_channels = channels\n        self.features = nn.Sequential(*features)\n\n        self.conv_last = nn.Sequential(\n            nn.Conv2d(p_channels, last_conv_channels, 1, 1, 0, bias=False),\n            nn.BatchNorm2d(last_conv_channels, affine=False),\n            nn.ReLU(inplace=True),\n        )\n        self.globalpool = nn.AvgPool2d(self._feature_map_size)\n        self.dropout = nn.Dropout(0.1)\n        self.classifier = nn.Sequential(\n            nn.Linear(last_conv_channels, n_classes, bias=False),\n        )\n\n        self._initialize_weights()\n\n    def _make_blocks(self, blocks, in_channels, channels):\n        result = []\n        for i in range(blocks):\n            stride = 2 if i == 0 else 1\n            inp = in_channels if i == 0 else channels\n            oup = channels\n\n            base_mid_channels = channels // 2\n            mid_channels = int(base_mid_channels)  # prepare for scale\n            choice_block = mutables.LayerChoice([\n                ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=3, stride=stride),\n                ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=5, stride=stride),\n                ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=7, stride=stride),\n                ShuffleXceptionBlock(inp, oup, mid_channels=mid_channels, stride=stride)\n            ])\n            result.append(choice_block)\n\n            # find the corresponding flops\n            flop_key = (inp, oup, mid_channels, self._feature_map_size, self._feature_map_size, stride)\n            self._parsed_flops[choice_block.key] = [\n                self._op_flops_dict[\"{}_stride_{}\".format(k, stride)][flop_key] for k in self.block_keys\n            ]\n            if stride == 2:\n                self._feature_map_size //= 2\n        return result\n\n    def forward(self, x):\n        bs = x.size(0)\n        x = self.first_conv(x)\n        x = self.features(x)\n        x = self.conv_last(x)\n        x = self.globalpool(x)\n\n        x = self.dropout(x)\n        x = x.contiguous().view(bs, -1)\n        x = self.classifier(x)\n        return x\n\n    def get_candidate_flops(self, candidate):\n        conv1_flops = self._op_flops_dict[\"conv1\"][(3, self._first_conv_channels,\n                                                    self._input_size, self._input_size, 2)]\n        # Should use `last_conv_channels` here, but megvii insists that it's `n_classes`. Keeping it.\n        # https://github.com/megvii-model/SinglePathOneShot/blob/36eed6cf083497ffa9cfe7b8da25bb0b6ba5a452/src/Supernet/flops.py#L313\n        rest_flops = self._op_flops_dict[\"rest_operation\"][(self.stage_channels[-1], self._n_classes,\n                                                            self._feature_map_size, self._feature_map_size, 1)]\n        total_flops = conv1_flops + rest_flops\n        for k, m in candidate.items():\n            parsed_flops_dict = self._parsed_flops[k]\n            if isinstance(m, dict):  # to be compatible with classical nas format\n                total_flops += parsed_flops_dict[m[\"_idx\"]]\n            else:\n                total_flops += parsed_flops_dict[torch.max(m, 0)[1]]\n        return total_flops\n\n    def _initialize_weights(self):\n        for name, m in self.named_modules():\n            if isinstance(m, nn.Conv2d):\n                if 'first' in name:\n                    nn.init.normal_(m.weight, 0, 0.01)\n                else:\n                    nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n            elif isinstance(m, nn.BatchNorm2d):\n                if m.weight is not None:\n                    nn.init.constant_(m.weight, 1)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0.0001)\n                nn.init.constant_(m.running_mean, 0)\n            elif isinstance(m, nn.BatchNorm1d):\n                nn.init.constant_(m.weight, 1)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0.0001)\n                nn.init.constant_(m.running_mean, 0)\n            elif isinstance(m, nn.Linear):\n                nn.init.normal_(m.weight, 0, 0.01)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n\n\ndef load_and_parse_state_dict(filepath=\"./data/checkpoint-150000.pth.tar\"):\n    checkpoint = torch.load(filepath, map_location=torch.device(\"cpu\"))\n    result = dict()\n    for k, v in checkpoint[\"state_dict\"].items():\n        if k.startswith(\"module.\"):\n            k = k[len(\"module.\"):]\n        result[k] = v\n    return result\n", "path": "examples/nas/spos/network.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# Licensed under the MIT license.\n\nimport os\nimport pickle\nimport re\n\nimport torch\nimport torch.nn as nn\nfrom nni.nas.pytorch import mutables\n\nfrom blocks import ShuffleNetBlock, ShuffleXceptionBlock\n\n\nclass ShuffleNetV2OneShot(nn.Module):\n    block_keys = [\n        'shufflenet_3x3',\n        'shufflenet_5x5',\n        'shufflenet_7x7',\n        'xception_3x3',\n    ]\n\n    def __init__(self, input_size=224, first_conv_channels=16, last_conv_channels=1024, n_classes=1000,\n                 op_flops_path=\"./data/op_flops_dict.pkl\"):\n        super().__init__()\n\n        assert input_size % 32 == 0\n        with open(os.path.join(os.path.dirname(__file__), op_flops_path), \"rb\") as fp:\n            self._op_flops_dict = pickle.load(fp)\n\n        self.stage_blocks = [4, 4, 8, 4]\n        self.stage_channels = [64, 160, 320, 640]\n        self._parsed_flops = dict()\n        self._input_size = input_size\n        self._feature_map_size = input_size\n        self._first_conv_channels = first_conv_channels\n        self._last_conv_channels = last_conv_channels\n        self._n_classes = n_classes\n\n        # building first layer\n        self.first_conv = nn.Sequential(\n            nn.Conv2d(3, first_conv_channels, 3, 2, 1, bias=False),\n            nn.BatchNorm2d(first_conv_channels, affine=False),\n            nn.ReLU(inplace=True),\n        )\n        self._feature_map_size //= 2\n\n        p_channels = first_conv_channels\n        features = []\n        for num_blocks, channels in zip(self.stage_blocks, self.stage_channels):\n            features.extend(self._make_blocks(num_blocks, p_channels, channels))\n            p_channels = channels\n        self.features = nn.Sequential(*features)\n\n        self.conv_last = nn.Sequential(\n            nn.Conv2d(p_channels, last_conv_channels, 1, 1, 0, bias=False),\n            nn.BatchNorm2d(last_conv_channels, affine=False),\n            nn.ReLU(inplace=True),\n        )\n        self.globalpool = nn.AvgPool2d(self._feature_map_size)\n        self.dropout = nn.Dropout(0.1)\n        self.classifier = nn.Sequential(\n            nn.Linear(last_conv_channels, n_classes, bias=False),\n        )\n\n        self._initialize_weights()\n\n    def _make_blocks(self, blocks, in_channels, channels):\n        result = []\n        for i in range(blocks):\n            stride = 2 if i == 0 else 1\n            inp = in_channels if i == 0 else channels\n            oup = channels\n\n            base_mid_channels = channels // 2\n            mid_channels = int(base_mid_channels)  # prepare for scale\n            choice_block = mutables.LayerChoice([\n                ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=3, stride=stride),\n                ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=5, stride=stride),\n                ShuffleNetBlock(inp, oup, mid_channels=mid_channels, ksize=7, stride=stride),\n                ShuffleXceptionBlock(inp, oup, mid_channels=mid_channels, stride=stride)\n            ])\n            result.append(choice_block)\n\n            # find the corresponding flops\n            flop_key = (inp, oup, mid_channels, self._feature_map_size, self._feature_map_size, stride)\n            self._parsed_flops[choice_block.key] = [\n                self._op_flops_dict[\"{}_stride_{}\".format(k, stride)][flop_key] for k in self.block_keys\n            ]\n            if stride == 2:\n                self._feature_map_size //= 2\n        return result\n\n    def forward(self, x):\n        bs = x.size(0)\n        x = self.first_conv(x)\n        x = self.features(x)\n        x = self.conv_last(x)\n        x = self.globalpool(x)\n\n        x = self.dropout(x)\n        x = x.contiguous().view(bs, -1)\n        x = self.classifier(x)\n        return x\n\n    def get_candidate_flops(self, candidate):\n        conv1_flops = self._op_flops_dict[\"conv1\"][(3, self._first_conv_channels,\n                                                    self._input_size, self._input_size, 2)]\n        # Should use `last_conv_channels` here, but megvii insists that it's `n_classes`. Keeping it.\n        # https://github.com/megvii-model/SinglePathOneShot/blob/36eed6cf083497ffa9cfe7b8da25bb0b6ba5a452/src/Supernet/flops.py#L313\n        rest_flops = self._op_flops_dict[\"rest_operation\"][(self.stage_channels[-1], self._n_classes,\n                                                            self._feature_map_size, self._feature_map_size, 1)]\n        total_flops = conv1_flops + rest_flops\n        for k, m in candidate.items():\n            parsed_flops_dict = self._parsed_flops[k]\n            if isinstance(m, dict):  # to be compatible with classical nas format\n                total_flops += parsed_flops_dict[m[\"_idx\"]]\n            else:\n                total_flops += parsed_flops_dict[torch.max(m, 0)[1]]\n        return total_flops\n\n    def _initialize_weights(self):\n        for name, m in self.named_modules():\n            if isinstance(m, nn.Conv2d):\n                if 'first' in name:\n                    nn.init.normal_(m.weight, 0, 0.01)\n                else:\n                    nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n            elif isinstance(m, nn.BatchNorm2d):\n                if m.weight is not None:\n                    nn.init.constant_(m.weight, 1)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0.0001)\n                nn.init.constant_(m.running_mean, 0)\n            elif isinstance(m, nn.BatchNorm1d):\n                nn.init.constant_(m.weight, 1)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0.0001)\n                nn.init.constant_(m.running_mean, 0)\n            elif isinstance(m, nn.Linear):\n                nn.init.normal_(m.weight, 0, 0.01)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n\n\ndef load_and_parse_state_dict(filepath=\"./data/checkpoint-150000.pth.tar\"):\n    checkpoint = torch.load(filepath, map_location=torch.device(\"cpu\"))\n    if \"state_dict\" in checkpoint:\n        checkpoint = checkpoint[\"state_dict\"]\n    result = dict()\n    for k, v in checkpoint.items():\n        if k.startswith(\"module.\"):\n            k = k[len(\"module.\"):]\n        result[k] = v\n    return result\n", "path": "examples/nas/spos/network.py"}]}