Repository: Lightning-Universe/lightning-flash

Issue:
Setting a Sampler with Arguments is not working with ImageClassificationData factory methods
## 🐛 Bug

This is related to the discussion: https://github.com/PyTorchLightning/pytorch-lightning/discussions/10005

I tried exactly that with creating a `ImageClassificationData` with the `from_data_frame` factory:

```python
ImageClassificationData.from_data_frame(
            ...
            sampler=WeightedRandomSampler
        )
```

It seems that the behaviour in `data_module._train_dataloader()` function has changed and the sampler is initialized there:

![image](https://user-images.githubusercontent.com/19696233/155719933-00ce97ee-7561-4dea-a523-dfefeebcce15.png)

I also found this in the docs:
<img width="1129" alt="image" src="https://user-images.githubusercontent.com/19696233/155720322-7e8ffe0b-ded0-42b0-aaec-945005686825.png">


It is expected to no longer pass a sampler instance, but a sampler class. I am wondering how i can pass all the parameters of the **WeightedRandomSampler** to this piece of code. I was not able to do it. 

### To Reproduce

```python
ImageClassificationData.from_data_frame(
            ...
            sampler=WeightedRandomSampler
            weights=...      # <-- i was trying to forward the arguments as kwargs, but was not working
        )
```

### Expected behavior

As a user i would expect to again pass the Sampler instance to the Factory method. 

### Environment

 - OS (e.g., Linux): Ubuntu 18.04. in Docker
 - Python version: 3.8
 - PyTorch/Lightning/Flash Version (e.g., 1.10/1.5/0.7): 0.7.0
 - GPU models and configuration: GPU
 - Any other relevant information:



Discussion:

Resolving patch:
diff --git a/flash/core/data/data_module.py b/flash/core/data/data_module.py
--- a/flash/core/data/data_module.py
+++ b/flash/core/data/data_module.py
@@ -69,6 +69,37 @@ class DataModule(pl.LightningDataModule):
         num_workers: The number of workers to use for parallelized loading.
         sampler: A sampler following the :class:`~torch.utils.data.sampler.Sampler` type.
             Will be passed to the DataLoader for the training dataset. Defaults to None.
+
+    Examples
+    ________
+
+    .. testsetup::
+
+        >>> from flash import DataModule
+        >>> from flash.core.utilities.stages import RunningStage
+        >>> from torch.utils.data.sampler import SequentialSampler, WeightedRandomSampler
+        >>> class TestInput(Input):
+        ...     def train_load_data(self, _):
+        ...         return [(0, 1, 2, 3), (0, 1, 2, 3)]
+        >>> train_input = TestInput(RunningStage.TRAINING, [1])
+
+    You can provide the sampler to use for the train dataloader using the ``sampler`` argument.
+    The sampler can be a function or type that needs the dataset as an argument:
+
+    .. doctest::
+
+        >>> datamodule = DataModule(train_input, sampler=SequentialSampler, batch_size=1)
+        >>> print(datamodule.train_dataloader().sampler)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
+        <torch.utils.data.sampler.SequentialSampler object at ...>
+
+    Alternatively, you can pass a sampler instance:
+
+    .. doctest::
+
+        >>> datamodule = DataModule(train_input, sampler=WeightedRandomSampler([0.1, 0.5], 2), batch_size=1)
+        >>> print(datamodule.train_dataloader().sampler)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
+        <torch.utils.data.sampler.WeightedRandomSampler object at ...>
+
     """
 
     input_transform_cls = InputTransform
@@ -84,7 +115,7 @@ def __init__(
         val_split: Optional[float] = None,
         batch_size: Optional[int] = None,
         num_workers: int = 0,
-        sampler: Optional[Type[Sampler]] = None,
+        sampler: Optional[Union[Callable, Sampler, Type[Sampler]]] = None,
         pin_memory: bool = True,
         persistent_workers: bool = False,
     ) -> None:
@@ -206,8 +237,10 @@ def _train_dataloader(self) -> DataLoader:
         if self.sampler is None:
             sampler = None
             shuffle = not isinstance(train_ds, IterableDataset)
-        else:
+        elif callable(self.sampler):
             sampler = self.sampler(train_ds)
+        else:
+            sampler = self.sampler
 
         if isinstance(getattr(self, "trainer", None), pl.Trainer):
             dataloader = self.trainer.lightning_module.process_train_dataset(


Test patch:
diff --git a/tests/core/data/test_data_module.py b/tests/core/data/test_data_module.py
--- a/tests/core/data/test_data_module.py
+++ b/tests/core/data/test_data_module.py
@@ -419,22 +419,28 @@ def validation_step(self, batch, batch_idx):
     trainer.fit(model, datamodule=datamodule)
 
 
+@pytest.mark.parametrize("sampler, callable", [(mock.MagicMock(), True), (mock.NonCallableMock(), False)])
 @mock.patch("flash.core.data.data_module.DataLoader")
-def test_dataloaders_with_sampler(mock_dataloader):
-    mock_sampler = mock.MagicMock()
+def test_dataloaders_with_sampler(mock_dataloader, sampler, callable):
+    train_input = TestInput(RunningStage.TRAINING, [1])
     datamodule = DataModule(
-        TestInput(RunningStage.TRAINING, [1]),
+        train_input,
         TestInput(RunningStage.VALIDATING, [1]),
         TestInput(RunningStage.TESTING, [1]),
         batch_size=2,
         num_workers=0,
-        sampler=mock_sampler,
+        sampler=sampler,
     )
-    assert datamodule.sampler is mock_sampler
+
+    assert datamodule.sampler is sampler
     dl = datamodule.train_dataloader()
+
+    if callable:
+        sampler.assert_called_once_with(train_input)
+
     kwargs = mock_dataloader.call_args[1]
     assert "sampler" in kwargs
-    assert kwargs["sampler"] is mock_sampler.return_value
+    assert kwargs["sampler"] is (sampler.return_value if callable else sampler)
     for dl in [datamodule.val_dataloader(), datamodule.test_dataloader()]:
         kwargs = mock_dataloader.call_args[1]
         assert "sampler" not in kwargs