losses to use CPU

metrics/criteria/clip_loss.py

@@ -7,7 +7,7 @@ class CLIPLoss(torch.nn.Module):
 
     def __init__(self, opts):
         super(CLIPLoss, self).__init__()
-        self.model, self.preprocess = clip.load("ViT-B/32", device="cuda")
+        self.model, self.preprocess = clip.load("ViT-B/32", device="cpu")
         self.upsample = torch.nn.Upsample(scale_factor=7)
         self.avg_pool = torch.nn.AvgPool2d(kernel_size=opts.stylegan_size // 32)
 

metrics/criteria/id_loss.py

@@ -9,7 +9,7 @@ class IDLoss(nn.Module):
         super(IDLoss, self).__init__()
         print('Loading ResNet ArcFace')
         self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
-        self.facenet.load_state_dict(torch.load(opts.ir_se50_weights))
+        self.facenet.load_state_dict(torch.load(opts.ir_se50_weights),  map_location=torch.device('cpu'))
         self.pool = torch.nn.AdaptiveAvgPool2d((256, 256))
         self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
         self.facenet.eval()