Fix pin_memory_thread not exiting quickly (#23646)

Summary:
fixes #23642
Pull Request resolved: #23646

Differential Revision: D16600874

Pulled By: soumith

fbshipit-source-id: 50f0828d774a558d6f21e9dd21135906bd5be128

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:

torch/utils/data/dataloader.py

@@ -886,9 +886,13 @@ def _shutdown_workers(self):
  # 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()