imporved gan loss

loss/gan.py

@@ -1,6 +1,5 @@
 import torch.nn as nn
 import torch.nn.functional as F
-import torch
 
 
 class GANLoss(nn.Module):
@@ -11,7 +10,7 @@ class GANLoss(nn.Module):
         self.fake_label_val = fake_label_val
         self.loss_type = loss_type
 
-    def single_forward(self, prediction, target_is_real: bool, is_discriminator=False):
+    def forward(self, prediction, target_is_real: bool, is_discriminator=False):
         """
         gan loss forward
         :param prediction: network prediction
@@ -38,14 +37,3 @@ class GANLoss(nn.Module):
             return loss
         else:
             raise NotImplementedError(f'GAN type {self.loss_type} is not implemented.')
-
-    def forward(self, prediction, target_is_real: bool, is_discriminator=False):
-        if isinstance(prediction, torch.Tensor):
-            # origin
-            return self.single_forward(prediction, target_is_real, is_discriminator)
-        elif isinstance(prediction, list):
-            # for multi scale discriminator, e.g. MultiScaleDiscriminator
-            loss = 0
-            for p in prediction:
-                loss += self.single_forward(p[-1], target_is_real, is_discriminator)
-            return loss