torch.set_default_device(device)

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
'''