16march_5

frameworkeval.py

@@ -13,7 +13,7 @@ class DF(nn.Module):
         self.perceptual_weight = 0.25
         self.ssim_weight = 0.25
         self.idsim_weight = 0.25
-        self.device = "cuda"
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.vgg = vgg16(pretrained=True).features[:16].to(device).eval()
         self.facenet = InceptionResnetV1(pretrained='vggface2').to(device).eval()
         for param in self.facenet.parameters():