Repository: Lightning-AI/pytorch-lightning
File: tests/trainer/test_dataloaders.py (Python)

Code before:
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')
def test_dataloader_distributed_sampler_already_attached(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend when DistSampler already included on dataloader """

    class ModelWithDataLoaderDistributedSampler(EvalModelTemplate):

        def train_dataloader(self):
            dataloader = super().train_dataloader()
            dist_sampler = DistributedSampler(dataloader.dataset, shuffle=False)
            return DataLoader(
                dataset,
                batch_size=self.batch_size,
                drop_last=False,
                sampler=dist_sampler,
                shuffle=False
            )

    model = ModelWithDataLoaderDistributedSampler()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=100,
        callbacks=[DistribSamplerCallback()]
        replace_sampler_ddp=True,
    )
    trainer.fit(model)
    assert result == 1, "DDP Training failed"



@pytest.mark.skipif(torch.cuda.device_count() < 3, reason='Test requires multiple GPUs')
def test_batch_size_smaller_than_num_gpus(tmpdir):
    # we need at least 3 gpus for this test
    num_gpus = 3
    batch_size = 3

    class CurrentTestModel(EvalModelTemplate):

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            # batch norm doesn't work with batch size 1, we replace it
            self.c_d1_bn = torch.nn.ReLU()

        def training_step(self, *args, **kwargs):
            output = super().training_step(*args, **kwargs)
            loss = output['loss']
            # we make sure to add some metrics to the output dict,
            # this is essential for this test
            output['progress_bar'] = {'train_loss': loss}

Review comment (suggestion):
```suggestion
        callbacks=[DistribSamplerCallback()],
```

Code after:

    def train_dataloader(self):
        dataloader = super().train_dataloader()
        dist_sampler = DistributedSampler(dataloader.dataset, shuffle=True)
        return DataLoader(
            dataloader.dataset,
            batch_size=self.batch_size,
            drop_last=False,
            sampler=dist_sampler,
            shuffle=False
        )


@pytest.mark.skipif(platform.system() == 'Windows', reason='Does not apply to Windows platform.')
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')
def test_dataloader_distributed_sampler_already_attached(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend when DistSampler already included on dataloader """

    model = ModelWithDataLoaderDistributedSampler()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=100,
        callbacks=[DistribSamplerCallback()],
        replace_sampler_ddp=True,
    )
    result = trainer.fit(model)
    assert result == 1, "DDP Training failed"


@pytest.mark.skipif(torch.cuda.device_count() < 3, reason='Test requires multiple GPUs')
def test_batch_size_smaller_than_num_gpus(tmpdir):
    # we need at least 3 gpus for this test
    num_gpus = 3
    batch_size = 3

    class CurrentTestModel(EvalModelTemplate):

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            # batch norm doesn't work with batch size 1, we replace it
            self.c_d1_bn = torch.nn.ReLU()

        def training_step(self, *args, **kwargs):
            output = super().training_step(*args, **kwargs)
            loss = output['loss']
            # we make sure to add some metrics to the output dict,
            # this is essential for this test
            output['progress_bar'] = {'train_loss': loss}