Update lightning_module_template.py

This commit is contained in:
William Falcon
2019-08-23 02:39:05 -04:00
committed by GitHub
parent 9fc66026f1
commit d5d47eab0d
@@ -194,7 +194,7 @@ class LightningTemplateModel(LightningModule):
batch_size = self.hparams.batch_size
try:
if self.on_gpu:
if self.trainer.use_ddp:
train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
batch_size = batch_size // self.trainer.world_size # scale batch size
except Exception: