fix device error

.idea/workspace.xml

@@ -6,8 +6,7 @@
   <component name="ChangeListManager">
     <list default="true" id="9fb9e207-fc4f-4ff3-9adc-3c4c1e67daa7" name="Changes" comment="">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/networks/_deeplab.py" beforeDir="false" afterPath="$PROJECT_DIR$/networks/_deeplab.py" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/networks/backbone/__init__.py" beforeDir="false" afterPath="$PROJECT_DIR$/networks/backbone/__init__.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/utils.py" beforeDir="false" afterPath="$PROJECT_DIR$/utils.py" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -49,7 +48,7 @@
       <option name="presentableId" value="Default" />
       <updated>1664204268713</updated>
       <workItem from="1664204270261" duration="37000" />
-      <workItem from="1664204316867" duration="5130000" />
+      <workItem from="1664204316867" duration="5530000" />
     </task>
     <servers />
   </component>

utils.py

@@ -7,8 +7,10 @@ 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)
+
     state_dict = torch.hub.load_state_dict_from_url(
-        "https://www.robots.ox.ac.uk/~vgg/research/namedmask/shared_files/voc2012/namedmask_voc2012.pt"
+        "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)