diff --git a/examples/new_project_templates/lightning_module_template.py b/examples/new_project_templates/lightning_module_template.py index 36b69812..826b6425 100644 --- a/examples/new_project_templates/lightning_module_template.py +++ b/examples/new_project_templates/lightning_module_template.py @@ -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