From 25dbd7a936050fe3277eac810d1038e23e3631ef Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Mon, 8 Jul 2019 19:45:52 -0400
Subject: [PATCH] scaled batch size

---
 examples/new_project_templates/lightning_module_template.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/new_project_templates/lightning_module_template.py b/examples/new_project_templates/lightning_module_template.py
index 19e5e0bd..b44e3de1 100644
--- a/examples/new_project_templates/lightning_module_template.py
+++ b/examples/new_project_templates/lightning_module_template.py
@@ -163,7 +163,7 @@ class LightningTemplateModel(LightningModule):
         batch_size = self.hparams.batch_size
 
         try:
-            if self.hparams.nb_gpu_nodes > 1:
+            if self.on_gpu:
                 train_sampler = DistributedSampler(dataset, num_replicas=self.trainer.world_size, rank=self.trainer.proc_rank)
 
                 # scale batch size