fix pin_memory_thread not exiting quickly (#23647)

bb4ff00f · Tongzhou Wang · Soumith Chintala · 49e32ffc · bb4ff00f · bb4ff00f
--- a/torch/utils/data/_utils/pin_memory.py
+++ b/torch/utils/data/_utils/pin_memory.py
@@ -22,7 +22,7 @@ def _pin_memory_loop(in_queue, out_queue, device_id, done_event):
        except queue.Empty:
            continue
        idx, data = r
-        if not isinstance(data, ExceptionWrapper):
+        if not done_event.is_set() and not isinstance(data, ExceptionWrapper):
            try:
                data = pin_memory(data)
            except Exception:

--- a/torch/utils/data/dataloader.py
+++ b/torch/utils/data/dataloader.py
@@ -886,9 +886,13 @@ class _MultiProcessingDataLoaderIter(_BaseDataLoaderIter):
                # corrupted data in `worker_result_queue` which `pin_memory_thread`
                # reads from.
                if hasattr(self, 'pin_memory_thread'):
-                    self.pin_memory_thread_done_event.set()
                    # Use hasattr in case error happens before we set the attribute.
+                    self.pin_memory_thread_done_event.set()
+                    # Send something to pin_memory_thread in case it is waiting
+                    # so that it can wake up and check `pin_memory_thread_done_event`
+                    self.worker_result_queue.put((None, None))
                    self.pin_memory_thread.join()
+                    self.worker_result_queue.close()

                # Exit workers now.
                self.workers_done_event.set()