Fix failure when DataLoader(batch_size=None) is passed (#10345)

* add test, + add change to data loading batch sample method

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Refactor and CHANGELOG

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Carlos Mocholí <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -77,7 +77,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed `apply_to_collection(defaultdict)` ([#10316](https://github.com/PyTorchLightning/pytorch-lightning/issues/10316))
 
 
--
+- Fixed failure when `DataLoader(batch_size=None)` is passed ([#10345](https://github.com/PyTorchLightning/pytorch-lightning/issues/10345))
 
 
 -

pytorch_lightning/trainer/data_loading.py

@@ -184,7 +184,7 @@ def _dataloader_init_kwargs_resolve_sampler(
         batch_sampler = getattr(dataloader, "batch_sampler")
         is_predicting = mode == RunningStage.PREDICTING
         # checking the batch sampler type is different than PyTorch default.
-        if (batch_sampler is not None and type(batch_sampler) is not BatchSampler) or is_predicting:
+        if batch_sampler is not None and (type(batch_sampler) is not BatchSampler or is_predicting):
             batch_sampler = type(batch_sampler)(
                 sampler,
                 batch_size=batch_sampler.batch_size,

tests/trainer/test_data_loading.py

@@ -283,25 +283,26 @@ class CustomSampler(Sampler):
         trainer.prepare_dataloader(dataloader, shuffle=True)
 
 
-def test_loader_detaching():
-    """Checks that the loader has been resetted after the entrypoint."""
+class LoaderTestModel(BoringModel):
+    def training_step(self, batch, batch_idx):
+        assert len(self.trainer.train_dataloader.loaders) == 10
+        return super().training_step(batch, batch_idx)
 
-    class LoaderTestModel(BoringModel):
-        def training_step(self, batch, batch_idx):
-            assert len(self.trainer.train_dataloader.loaders) == 10
-            return super().training_step(batch, batch_idx)
+    def validation_step(self, batch, batch_idx):
+        assert len(self.trainer.val_dataloaders[0]) == 10
+        return super().validation_step(batch, batch_idx)
 
-        def validation_step(self, batch, batch_idx):
-            assert len(self.trainer.val_dataloaders[0]) == 10
-            return super().validation_step(batch, batch_idx)
+    def test_step(self, batch, batch_idx):
+        assert len(self.trainer.test_dataloaders[0]) == 10
+        return super().test_step(batch, batch_idx)
 
-        def test_step(self, batch, batch_idx):
-            assert len(self.trainer.test_dataloaders[0]) == 10
-            return super().test_step(batch, batch_idx)
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        assert len(self.trainer.predict_dataloaders[0]) == 10
+        return super().predict_step(batch, batch_idx, dataloader_idx=dataloader_idx)
 
-        def predict_step(self, batch, batch_idx, dataloader_idx=0):
-            assert len(self.trainer.predict_dataloaders[0]) == 10
-            return super().predict_step(batch, batch_idx, dataloader_idx=dataloader_idx)
+
+def test_loader_detaching():
+    """Checks that the loader has been resetted after the entrypoint."""
 
     loader = DataLoader(RandomDataset(32, 10), batch_size=1)
 
@@ -340,3 +341,10 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0):
     assert len(model.val_dataloader()) == 64
     assert len(model.predict_dataloader()) == 64
     assert len(model.test_dataloader()) == 64
+
+
+def test_pre_made_batches():
+    """Check that loader works with pre-made batches."""
+    loader = DataLoader(RandomDataset(32, 10), batch_size=None)
+    trainer = Trainer(fast_dev_run=1)
+    trainer.predict(LoaderTestModel(), loader)