9 lines
229 B
Python
9 lines
229 B
Python
import torch
|
|
|
|
def get_default_provider():
|
|
if torch.cuda.is_available():
|
|
return "CUDA"
|
|
elif hasattr(torch, "backends") and hasattr(torch.backends, "rocm"):
|
|
return "ROCM"
|
|
else:
|
|
return "CPU" |