updated args

docs/source/examples/example_model.py

@@ -78,7 +78,7 @@ class ExampleModel(RootModule):
         output = OrderedDict({
             'loss_val': loss_val,
         })
-        return torch.tensor(4)
+        return output
 
     def validation_step(self, data_batch, batch_i):
         """
@@ -102,7 +102,7 @@ class ExampleModel(RootModule):
             'loss_val': loss_val,
             'val_acc': val_acc,
         })
-        return torch.tensor(4)
+        return output
 
 
     def validation_end(self, outputs):

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -72,9 +72,9 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
                 # ---------------
                 # CHANGE
                 if module.training:
-                    return module.training_step(*input, **kwargs)
+                    output = module.training_step(*input, **kwargs)
                 else:
-                    return module.validation_step(*input, **kwargs)
+                    output = module.validation_step(*input, **kwargs)
                 # ---------------
 
             with lock: