(self, batch, gpu_id)
| 790 | m.single_gpu = self.single_gpu |
| 791 | |
| 792 | def transfer_batch_to_gpu(self, batch, gpu_id): |
| 793 | # base case: object can be directly moved using `cuda` or `to` |
| 794 | if callable(getattr(batch, 'cuda', None)): |
| 795 | return batch.cuda(gpu_id, non_blocking=True) |
| 796 | |
| 797 | elif callable(getattr(batch, 'to', None)): |
| 798 | return batch.to(torch.device('cuda', gpu_id), non_blocking=True) |
| 799 | |
| 800 | # when list |
| 801 | elif isinstance(batch, list): |
| 802 | for i, x in enumerate(batch): |
| 803 | batch[i] = self.transfer_batch_to_gpu(x, gpu_id) |
| 804 | return batch |
| 805 | |
| 806 | # when tuple |
| 807 | elif isinstance(batch, tuple): |
| 808 | batch = list(batch) |
| 809 | for i, x in enumerate(batch): |
| 810 | batch[i] = self.transfer_batch_to_gpu(x, gpu_id) |
| 811 | return tuple(batch) |
| 812 | |
| 813 | # when dict |
| 814 | elif isinstance(batch, dict): |
| 815 | for k, v in batch.items(): |
| 816 | batch[k] = self.transfer_batch_to_gpu(v, gpu_id) |
| 817 | |
| 818 | return batch |
| 819 | |
| 820 | # nothing matches, return the value as is without transform |
| 821 | return batch |
| 822 | |
| 823 | def set_distributed_mode(self, distributed_backend): |
| 824 | # skip for CPU |
no outgoing calls
no test coverage detected