scaled batch size

This commit is contained in:
William Falcon
2019-07-08 19:54:00 -04:00
parent f95fad864d
commit a87073bffd
@@ -164,6 +164,8 @@ class LightningTemplateModel(LightningModule):
# try:
# if self.on_gpu:
import torch.distributed as dist
print(dist.get_world_size())
train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
batch_size = batch_size // self.trainer.world_size # scale batch size