scaled batch size

examples/new_project_templates/lightning_module_template.py

@@ -162,16 +162,12 @@ class LightningTemplateModel(LightningModule):
         train_sampler = None
         batch_size = self.hparams.batch_size
 
-        # try:
-        #     if self.on_gpu:
-        import torch.distributed as dist
-        print(dist.get_world_size())
+        try:
+            if self.on_gpu:
                 train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
                 batch_size = batch_size // self.trainer.world_size  # scale batch size
-        print(batch_size)
-
-        # except Exception as e:
-        #     pass
+        except Exception as e:
+            pass
 
         should_shuffle = train_sampler is None
         loader = DataLoader(