Add optional gan loss to MSE

htem · Nov 20, 2023 · bb54f1d · bb54f1d
1 parent 546457a
commit bb54f1d
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/autoseg/losses/mtlsd_losses.py b/src/autoseg/losses/mtlsd_losses.py
@@ -40,9 +40,11 @@ def forward(
         aff_loss = self.aff_lambda * self._calc_loss(pred_affs, gt_affs, affs_weights)
 
         # calculate MSE loss for GAN errors
-        real_scores = self.discriminator(gt_enhanced)
-        fake_scores = self.discriminator(pred_enhanced)
-
-        gan_loss = self.gan_lambda * (torch.mean((real_scores - 1) ** 2) + torch.mean(fake_scores ** 2))
+        if gt_enhanced is not None and pred_enhanced is not None:
+            real_scores = self.discriminator(gt_enhanced)
+            fake_scores = self.discriminator(pred_enhanced)
+            gan_loss = self.gan_lambda * (torch.mean((real_scores - 1) ** 2) + torch.mean(fake_scores ** 2))
+        else:
+            gan_loss: float = 0.
 
         return lsd_loss + aff_loss + gan_loss