Update modules/frameworkeval.py

modules/frameworkeval.py

@@ -14,8 +14,8 @@ class DF(nn.Module):
         self.ssim_weight = 0.25
         self.idsim_weight = 0.25
         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()
+        self.vgg = vgg16(pretrained=True).features[:16].to(self.device).eval()
+        self.facenet = InceptionResnetV1(pretrained='vggface2').to(self.device).eval()
         for param in self.facenet.parameters():
             param.requires_grad = False  # Freeze the model
         self.cosloss = nn.CosineEmbeddingLoss()
@@ -29,10 +29,10 @@ class DF(nn.Module):
     def idsimilarity(self, real, fake):
         with torch.no_grad():
             # Extract embeddings
-            input_embed = self.facenet(real).to(device)
-            generated_embed = self.facenet(fake).to(device)
+            input_embed = self.facenet(real).to(self.device)
+            generated_embed = self.facenet(fake).to(self.device)
         # Compute cosine similarity loss
-        target = torch.ones(input_embed.size(0)).to(real.device)  # Target = 1 (maximize similarity)
+        target = torch.ones(input_embed.size(0)).to(self.device)  # Target = 1 (maximize similarity)
         return self.cosloss(input_embed, generated_embed, target)
 
     def forward(self, r, f):