| import torch | |
| def get_device(device_arg): | |
| if device_arg == "auto": | |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| return torch.device(device_arg) | |
| import torch | |
| def get_device(device_arg): | |
| if device_arg == "auto": | |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| return torch.device(device_arg) | |