From 440f47b864e51b079f2a7644f9d89990f78ebc53 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Tue, 25 Jun 2019 20:24:03 -0400
Subject: [PATCH] updated args

---
 pytorch_lightning/models/trainer.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index f5264b79..5ab2e374 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -196,7 +196,10 @@ class Trainer(TrainerIO):
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
-            output = model(data_batch, batch_i)
+            if self.data_parallel:
+                output = model(data_batch, batch_i)
+            else:
+                output = model.validation_step(data_batch, batch_i)
 
             # when DP, we need to aggregate the scalars we received as outputs
             # use mean as the reduce function