File size: 451 Bytes
3ea3da7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#!/usr/bin/env python3

import sys

try:
    import torch
except ImportError:
    print("Torch is required for configuring a kernel build.", file=sys.stderr)
    sys.exit(1)

if torch.version.cuda is not None:
    print("CUDA")
elif torch.version.hip is not None:
    print("HIP")
elif torch.backends.mps.is_available():
    print("METAL")
elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
    print("SYCL")
else:
    print("CPU")