fix state dict loading error

.idea/workspace.xml

@@ -48,7 +48,7 @@
       <option name="presentableId" value="Default" />
       <updated>1664204268713</updated>
       <workItem from="1664204270261" duration="37000" />
-      <workItem from="1664204316867" duration="5530000" />
+      <workItem from="1664204316867" duration="5840000" />
     </task>
     <servers />
   </component>

utils.py

@@ -8,13 +8,14 @@ from networks import convert_to_separable_conv, set_bn_momentum
 def get_network() -> torch.nn.Module:
     network = deeplabv3plus_resnet50(num_classes=21, pretrained_backbone=False)
 
+    convert_to_separable_conv(network.classifier)
+    set_bn_momentum(network.backbone, momentum=0.01)
+
     state_dict = torch.hub.load_state_dict_from_url(
         "https://www.robots.ox.ac.uk/~vgg/research/namedmask/shared_files/voc2012/namedmask_voc2012.pt",
         map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu")
     )
-    network.backbone.load_state_dict(state_dict, strict=True)
-    convert_to_separable_conv(network.classifier)
-    set_bn_momentum(network.backbone, momentum=0.01)
+    network.load_state_dict(state_dict, strict=True)
     return network