update type

Lightning-AI · Nov 19, 2021 · 24a85ce · 24a85ce
1 parent f6b11ed
commit 24a85ce
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_lightning/lite/lite.py b/pytorch_lightning/lite/lite.py
@@ -188,7 +188,7 @@ def setup(
 
     def setup_dataloaders(
         self, *dataloaders: DataLoader, replace_sampler: bool = True, move_to_device: bool = True
-    ) -> Union[Iterable, List[Iterable]]:
+    ) -> Union[_LiteDataLoader, List[_LiteDataLoader]]:
         """Setup one or multiple dataloaders for accelerated training. If you need different settings for each
         dataloader, call this method individually for each one.
 
@@ -213,7 +213,7 @@ def setup_dataloaders(
 
     def _setup_dataloader(
         self, dataloader: DataLoader, replace_sampler: bool = True, move_to_device: bool = True
-    ) -> Iterable:
+    ) -> _LiteDataLoader:
         """Setup a single dataloader for accelerated training.
 
         Args: