mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-10-27 13:12:20 +08:00
scaled batch size
This commit is contained in:
1 parent
f95fad864d
commit
a87073bffd
1 file changed
+2
@@ -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
|
||||
|
||||
|
||||
Reference in new issue
Block a user