| 542 | self.initialized_events = initialized_events |
| 543 | |
| 544 | def transfer_data_to_cuda(self, data, device): |
| 545 | if data is None: |
| 546 | return None |
| 547 | else: |
| 548 | if isinstance(data, torch.Tensor): |
| 549 | data = data.to(device) |
| 550 | elif isinstance(data, list): |
| 551 | data = [ |
| 552 | self.transfer_data_to_cuda(subdata, device) |
| 553 | for subdata in data |
| 554 | ] |
| 555 | elif isinstance(data, dict): |
| 556 | data = { |
| 557 | key: self.transfer_data_to_cuda(val, device) |
| 558 | for key, val in data.items() |
| 559 | } |
| 560 | return data |
| 561 | |
| 562 | def mp_worker(self, gpu, gpu_infer, pmi_rank, pmi_world_size, in_q_list, |
| 563 | out_q, initialized_events, work_env): |