(in_queue, out_queue, done_event, pin_memory, device_id)
| 65 | |
| 66 | |
| 67 | def _worker_manager_loop(in_queue, out_queue, done_event, pin_memory, device_id): |
| 68 | if pin_memory: |
| 69 | torch.cuda.set_device(device_id) |
| 70 | |
| 71 | while True: |
| 72 | try: |
| 73 | r = in_queue.get() |
| 74 | except Exception: |
| 75 | if done_event.is_set(): |
| 76 | return |
| 77 | raise |
| 78 | if r is None: |
| 79 | break |
| 80 | if isinstance(r[1], ExceptionWrapper): |
| 81 | out_queue.put(r) |
| 82 | continue |
| 83 | idx, batch = r |
| 84 | try: |
| 85 | if pin_memory: |
| 86 | batch = pin_memory_batch(batch) |
| 87 | except Exception: |
| 88 | out_queue.put((idx, ExceptionWrapper(sys.exc_info()))) |
| 89 | else: |
| 90 | out_queue.put((idx, batch)) |
| 91 | |
| 92 | numpy_type_map = { |
| 93 | 'float64': torch.DoubleTensor, |
nothing calls this directly
no test coverage detected