improvements

Lightning-AI · Jan 19, 2022 · 108fcfa · 108fcfa
1 parent 0e0ebd8
commit 108fcfa
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/pytorch_lightning/trainer/connectors/data_connector.py b/pytorch_lightning/trainer/connectors/data_connector.py
@@ -397,7 +397,8 @@ def _resolve_sampler(self, dataloader: DataLoader, shuffle: bool, mode: Optional
                     f"Using `DistributedSampler` with the dataloaders. During trainer.{trainer_fn.value}(),"
                     " it is recommended to use `Trainer(devices=1)` to ensure each sample/batch gets evaluated"
                     " exactly once. Otherwise, multi-device settings use `DistributedSampler` that replicates"
-                    " some samples to make sure all devices have same batch size in case of uneven inputs."
+                    " some samples to make sure all devices have same batch size in case of uneven inputs.",
+                    category=PossibleUserWarning,
                 )
 
             return sampler

diff --git a/tests/trainer/connectors/test_data_connector.py b/tests/trainer/connectors/test_data_connector.py
@@ -19,6 +19,7 @@
 from pytorch_lightning import Trainer
 from pytorch_lightning.trainer.connectors.data_connector import _DataLoaderSource
 from pytorch_lightning.trainer.states import TrainerFn
+from pytorch_lightning.utilities.warnings import PossibleUserWarning
 from tests.helpers import BoringDataModule, BoringModel
 
 
@@ -77,9 +78,9 @@ def test_eval_distributed_sampler_warning(tmpdir):
     trainer._data_connector.attach_data(model)
 
     trainer.state.fn = TrainerFn.VALIDATING
-    with pytest.warns(UserWarning, match="multi-device settings use `DistributedSampler`"):
+    with pytest.warns(PossibleUserWarning, match="multi-device settings use `DistributedSampler`"):
         trainer.reset_val_dataloader(model)
 
     trainer.state.fn = TrainerFn.TESTING
-    with pytest.warns(UserWarning, match="multi-device settings use `DistributedSampler`"):
+    with pytest.warns(PossibleUserWarning, match="multi-device settings use `DistributedSampler`"):
         trainer.reset_test_dataloader(model)