fix device

criteria/id_loss.py

@@ -17,7 +17,7 @@ class IDLoss(nn.Module):
     [4] https://github.com/eladrich/pixel2style2pixel
     """
 
-    def __init__(self, model_path, official=False):
+    def __init__(self, model_path, official=False, device="cpu"):
         """
         Arguments:
             model_path (str): Path to IR-SE50 model.
@@ -32,7 +32,7 @@ class IDLoss(nn.Module):
                 input_size=112, num_layers=50, drop_ratio=0.6, mode="ir_se"
             )
 
-        self.facenet.load_state_dict(torch.load(model_path))
+        self.facenet.load_state_dict(torch.load(model_path, map_location=device))
         self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
         self.facenet.eval()
 

criteria/mask.py

@@ -9,7 +9,7 @@ import numpy as np
 
 
 class Mask(nn.Module):
-    def __init__(self):
+    def __init__(self, device="cpu"):
         """
 
         |  Class     | Number | Class | Number |
@@ -41,7 +41,7 @@ class Mask(nn.Module):
             .requires_grad_(False)
         )
 
-        ckpt = torch.load(paths_config.deeplab, map_location=global_config.device)
+        ckpt = torch.load(paths_config.deeplab, map_location=device)
         state_dict = {
             k[7:]: v for k, v in ckpt["state_dict"].items() if "tracked" not in k
         }

training/coaches/base_coach.py

@@ -51,13 +51,14 @@ class BaseCoach:
             id_loss.IDLoss(
                 paths_config.ir_se50,
                 official=False,
+                device=global_config.device
             )
             .to(global_config.device)
             .eval()
         )
 
         if hyperparameters.use_mask:
-            self.mask = mask.Mask()
+            self.mask = mask.Mask(device=global_config.device)
 
         self.restart_training()