fix deeplab device

criteria/deeplab.py

@@ -309,7 +309,7 @@ def resnet50(pretrained=False, **kwargs):
     return model
 
 
-def resnet101(path=None, pretrained=False, num_groups=None, weight_std=False, **kwargs):
+def resnet101(path=None, pretrained=False, num_groups=None, weight_std=False, device="cpu", **kwargs):
     """Constructs a ResNet-101 model.
 
     Args:
@@ -326,7 +326,7 @@ def resnet101(path=None, pretrained=False, num_groups=None, weight_std=False, **
         model_dict = model.state_dict()
         if num_groups and weight_std:
             path = os.path.join(os.path.dirname(path), "R-101-GN-WS.pth.tar")
-            pretrained_dict = torch.load(path)
+            pretrained_dict = torch.load(path, map_location=device)
             overlap_dict = {
                 k[7:]: v for k, v in pretrained_dict.items() if k[7:] in model_dict
             }

criteria/mask.py

@@ -36,6 +36,7 @@ class Mask(nn.Module):
                 num_groups=32,
                 weight_std=True,
                 beta=False,
+                device=device,
             )
             .eval()
             .requires_grad_(False)