Fix dataloader._shutdown_workers if not all workers are started (#23762)

ssnl · soumith · commit 1ac19ea8bac1 · 2019-08-04T23:28:47.000-04:00
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
@@ -896,7 +896,10 @@ def _shutdown_workers(self):
 
                 # Exit workers now.
                 self.workers_done_event.set()
-                for worker_id in range(self.num_workers):
+                for worker_id in range(len(self.workers)):
+                    # Get number of workers from `len(self.workers)` instead of
+                    # `self.num_workers` in case we error before starting all
+                    # workers.
                     if self.workers_status[worker_id]:
                         self._shutdown_worker(worker_id)
                 for w in self.workers:
