set _prefetch_factor at _MultiProcessingDataLoaderIter and mypy fixes

blankstatic · pytorchmergebot · commit d6b83a767405 · 2022-11-17T20:28:38.000Z
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
@@ -245,7 +245,7 @@ def __init__(self, dataset: Dataset[T_co], batch_size: Optional[int] = 1,
                              'let num_workers > 0 to enable multiprocessing, otherwise set prefetch_factor to None.')
         elif num_workers > 0 and prefetch_factor is None:
             prefetch_factor = 2
-        elif prefetch_factor < 0:
+        elif prefetch_factor is not None and prefetch_factor < 0:
             raise ValueError('prefetch_factor option should be non-negative')
 
         if persistent_workers and num_workers == 0:
@@ -584,7 +584,6 @@ def __init__(self, loader: DataLoader) -> None:
         ws, rank = _get_distributed_settings()
         self._world_size = ws
         self._rank = rank
-        self._prefetch_factor = loader.prefetch_factor
         # for other backends, pin_memory_device need to set. if not set
         # default behaviour is CUDA device. if pin_memory_device is selected
         # and pin_memory is not set, the default behaviour false.
@@ -994,6 +993,8 @@ class _MultiProcessingDataLoaderIter(_BaseDataLoaderIter):
     def __init__(self, loader):
         super(_MultiProcessingDataLoaderIter, self).__init__(loader)
 
+        self._prefetch_factor = loader.prefetch_factor
+
         assert self._num_workers > 0
         assert self._prefetch_factor > 0
 
