inputs

Lightning-AI · Dec 18, 2020 · 7655c37 · 7655c37
1 parent 490acf2
commit 7655c37
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pytorch_lightning/overrides/data_parallel.py b/pytorch_lightning/overrides/data_parallel.py
@@ -159,13 +159,13 @@ def __init__(self, lightning_module: LightningModule):
 
     def forward(self, *inputs, **kwargs):
         if self.module.training:
-            output = self.module.training_step(*inputs[0], **kwargs[0])
+            output = self.module.training_step(*inputs, **kwargs)
             fx_called = 'training_step'
         elif self.module.testing:
-            output = self.module.test_step(*inputs[0], **kwargs[0])
+            output = self.module.test_step(*inputs, **kwargs)
             fx_called = 'test_step'
         else:
-            output = self.module.validation_step(*inputs[0], **kwargs[0])
+            output = self.module.validation_step(*inputs, **kwargs)
             fx_called = 'validation_step'
         if output is None:
             warn_missing_output(f'{fx_called} returned None. Did you forget to return an output?')