wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update lightning_module_template.py

Browse Source
This commit is contained in:
William Falcon
2019-08-23 02:42:40 -04:00
committed by GitHub
parent d5d47eab0d
commit c30f69f60d
1 changed files with 4 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/new_project_templates/lightning_module_template.py
+4 -7
View File
@@ -189,16 +189,13 @@ class LightningTemplateModel(LightningModule):
dataset = MNIST(root=self.hparams.data_root, train=train,
transform=transform, download=True)
# when using multi-node we need to add the datasampler
# when using multi-node (ddp) we need to add the datasampler
train_sampler = None
batch_size = self.hparams.batch_size
try:
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:
pass
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
should_shuffle = train_sampler is None
loader = DataLoader(