torch.set_default_device(device)
- device: 要设置为默认的设备
设置存储数据的默认设备（默认是cpu，训练模型时，用gpu，以提高训练速度）
import torch
# 打印当前所有可用的设备
# 检查CUDA是否可用
if torch.cuda.is_available():# 获取CUDA设备的数量n_gpu = torch.cuda.device_count()print(f"检测到 {n_gpu} 个CUDA设备。")# 遍历设备并打印它们的名称和索引for i in range(n_gpu):print(f"设备 {i}: {torch.cuda.get_device_name(i)}")
else:print("没有检测到CUDA设备。")# 打印当前默认设备
print(torch.get_default_device()) # cpu
# 设置默认设备为cuda:0
torch.set_default_device('cuda:0')
# torch.set_default_device('cuda')
# 再次打印当前默认设备
print(torch.get_default_device()) # cuda:0
'''output
检测到 1 个CUDA设备。
设备 0: NVIDIA GeForce GTX 1050 Ti
cpu
cuda:0
'''
)
