Update loss.py

torchgan · Nov 10, 2018 · d0516ca · d0516ca
1 parent c44509f
commit d0516ca
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchgan/losses/loss.py b/torchgan/losses/loss.py
@@ -59,7 +59,9 @@ def __init__(self, reduction='elementwise_mean', override_train_ops=None):
         self.reduction = reduction
         self.override_train_ops = override_train_ops
 
-    def train_ops(self, generator, discriminator, optimizer_discriminator, real_inputs, batch_size, device,
+    # NOTE(avik-pal): batch_size and device gets flipped if the order is not given as below. Investigate this
+    #                 error as might affect our support for custom loss functions.
+    def train_ops(self, generator, discriminator, optimizer_discriminator, real_inputs, device, batch_size,
                   labels=None):
         if self.override_train_ops is not None:
             return self.override_train_ops(self, generator, discriminator, optimizer_discriminator,