# cpu
device = torch.device("cpu")
# cuda
torch.cuda.set_device(device)
如何查看:
设置完之后,pdb打断点输出
print(data.device)
如果是cpu/cuda,会输出对应信息
# cpu
device = torch.device("cpu")
# cuda
torch.cuda.set_device(device)
如何查看:
设置完之后,pdb打断点输出
print(data.device)
如果是cpu/cuda,会输出对应信息