(device)
| 779 | |
| 780 | |
| 781 | def _get_device(device): |
| 782 | device = torch.device(device) |
| 783 | if device.type == "cuda" and device.index is None: |
| 784 | device = torch.device("cuda", torch.cuda.current_device()) |
| 785 | return device |
| 786 | |
| 787 | |
| 788 | class DataLoader(torch.utils.data.DataLoader): |