76d7a6e
1
2
3
4
5
6
7
8
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)
