Revert part of #10279

pytorch_lightning/lite/lite.py

@@ -238,18 +238,16 @@ def _setup_dataloader(
                 )
             sampler = self._get_distributed_sampler(dataloader, **self._strategy.distributed_sampler_kwargs)
 
+        # the dataloader needs to be re-instantiated because we want to update the input arguments (e.g., sampler)
         dataloader_kwargs = TrainerDataLoadingMixin._get_dataloader_init_kwargs(dataloader, sampler)
-        try:
-            dataloader = type(dataloader)(**dataloader_kwargs)
-        except TypeError:
-            dataloader_kwargs.pop("dataset")
-            dataloader = type(dataloader)(**dataloader_kwargs)
+        dataloader = type(dataloader)(**dataloader_kwargs)
+
         # add worker_init_fn for correct seeding in worker processes
         TrainerDataLoadingMixin._auto_add_worker_init_fn(dataloader, self.global_rank)
-        return _LiteDataLoader(
-            dataloader=self._strategy.process_dataloader(dataloader),
-            device=self.device if move_to_device and not isinstance(self._strategy, TPUSpawnPlugin) else None,
-        )
+
+        dataloader = self._strategy.process_dataloader(dataloader)
+        device = self.device if move_to_device and not isinstance(self._strategy, TPUSpawnPlugin) else None
+        return _LiteDataLoader(dataloader=dataloader, device=device)
 
     def backward(self, tensor: Tensor, *args: Any, model: Optional[_LiteModule] = None, **kwargs: Any) -> None:
         """Replaces ``loss.backward()`` in your training loop. Handles precision and automatically for you.

tests/lite/test_lite.py

@@ -192,57 +192,6 @@ def run(self):
     LiteWithCustomDataLoader().run()
 
 
-def test_setup_custom_dataloaders():
-    """Test that the setup_dataloaders method returns the dataloaders wrapped as LiteDataLoader."""
-    lite = EmptyLite()
-
-    class CustomDataLoader(DataLoader):
-        def __init__(self, value: int = 2, *args, **kwargs):
-            self.value = value
-            super().__init__(range(value), *args, **kwargs)
-
-    dataloader = CustomDataLoader(2, batch_size=2)
-
-    # single dataloader
-    lite_dataloader = lite.setup_dataloaders(dataloader)
-    assert lite_dataloader._dataloader
-    assert lite_dataloader.value == 2
-    batch0 = next(iter(lite_dataloader))
-    assert torch.equal(batch0, torch.tensor([0, 1]))
-
-    class CustomDataLoader2(DataLoader):
-        def __init__(self, range, *args, **kwargs):
-            self.range = range
-            super().__init__(range, *args, **kwargs)
-
-    dataloader = CustomDataLoader2(range(2), batch_size=2)
-
-    # single dataloader
-    lite_dataloader = lite.setup_dataloaders(dataloader)
-    assert lite_dataloader._dataloader
-    batch0 = next(iter(lite_dataloader))
-    assert torch.equal(batch0, torch.tensor([0, 1]))
-
-    class CustomDataLoader(DataLoader):
-        def __init__(self, value: int, *args, **kwargs):
-            super().__init__(range(value), *args, **kwargs)
-
-    class LiteWithCustomDataLoader(LightningLite):
-        def run(self):
-            # This doesn't fail as the context manager would save all the arguments provided
-            # to the dataloaders.
-            dataloader = CustomDataLoader(2, batch_size=2)
-            self.setup_dataloaders(dataloader)
-
-    LiteWithCustomDataLoader().run()
-
-    with pytest.raises(
-        MisconfigurationException, match="Trying to inject `DistributedSampler` into the `CustomDataLoader` instance"
-    ):
-        dataloader = CustomDataLoader(2, batch_size=2)
-        lite_dataloader = lite.setup_dataloaders(dataloader)
-
-
 def test_setup_dataloaders_twice_fails():
     """Test that calling setup_dataloaders with a dataloader that is already wrapped fails."""
     lite = EmptyLite()