18 lines
420 B
Python
18 lines
420 B
Python
import torch
|
||
|
||
|
||
def print_gpu_availability():
|
||
"""打印PyTorch的GPU可用性"""
|
||
if torch.cuda.is_available():
|
||
print(f"GPU可用：{torch.cuda.get_device_name(0)}")
|
||
else:
|
||
print("GPU不可用")
|
||
|
||
|
||
def get_gpu_device() -> torch.device:
|
||
"""获取PyTorch的GPU设备"""
|
||
if torch.cuda.is_available():
|
||
return torch.device("cuda")
|
||
else:
|
||
raise Exception("找不到CUDA！")
|