Add dataloader name in warning and wrap only dataloader_method in try

Author: kingyiusuen

pytorch_lightning/core/datamodule.py

@@ -493,11 +493,12 @@ def __len__(self) -> int:
         num_batches = 0
         not_implemented_count = 0
 
-        def get_num_batches(dataloader: DataLoader) -> None:
+        def get_num_batches(dataloader: DataLoader, name: str) -> None:
             nonlocal num_batches
             if not has_len(dataloader):
                 rank_zero_warn(
-                    "The number of batches for a dataloader is counted as 0 because it does not have `__len__` defined."
+                    f"The number of batches for a dataloader in `{name}` is counted as 0 "
+                    "because it does not have `__len__` defined."
                 )
             else:
                 num_batches += len(dataloader)
@@ -506,11 +507,12 @@ def get_num_batches(dataloader: DataLoader) -> None:
             dataloader_method = getattr(self, method_name)
             try:
                 dataloader = dataloader_method()
-                if isinstance(dataloader, CombinedLoader):
-                    dataloader = dataloader.loaders
-                apply_to_collection(dataloader, DataLoader, get_num_batches)
             except NotImplementedError:
                 not_implemented_count += 1
+                continue
+            if isinstance(dataloader, CombinedLoader):
+                dataloader = dataloader.loaders
+            apply_to_collection(dataloader, DataLoader, get_num_batches, method_name)
 
         if not_implemented_count == 4:
             rank_zero_warn("You datamodule does not have any valid dataloader so `__len__` will be returned as 0.")

tests/core/test_datamodules.py

@@ -664,7 +664,7 @@ def __len__(self):
     dataloader = CustomNotImplementedErrorDataloader(RandomDataset(1, 32))
     dm = LightningDataModule()
     setattr(dm, method_name, lambda: dataloader)
-    with pytest.warns(UserWarning, match="The number of batches for a dataloader is counted as 0"):
+    with pytest.warns(UserWarning, match=f"The number of batches for a dataloader in `{method_name}` is counted as 0"):
         assert len(dm) == 0