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

Code before:
        assert not test_sampler.shuffle


@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(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend """

    model = EvalModelTemplate()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=1,
        callbacks=[DistribSamplerCallback()]
    )
    trainer.fit(model)
    trainer.test(ckpt_path=None)


class ModelWithDataLoaderDistributedSampler(EvalModelTemplate):

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


@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)

Review comment (suggestion):
```suggestion
        dist_sampler = DistributedSampler(dataloader.dataset, shuffle=True)
```

Code after:
        assert not test_sampler.shuffle


@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(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend """

    model = EvalModelTemplate()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=1,
        callbacks=[DistribSamplerCallback()]
    )
    trainer.fit(model)
    trainer.test(ckpt_path=None)


class ModelWithDataLoaderDistributedSampler(EvalModelTemplate):

    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)