zl程序教程

您现在的位置是:首页 >  硬件

当前栏目

判断torch.device是gpu还是cpu

CPU 判断 还是 GPU device torch
2023-09-14 09:05:44 时间

判断torch.device是gpu还是cpu

这样直接判断是有问题,会有问题,因为torch.device是class类型

 device=torch.device('cuda:0')
 if device!= "cpu":
    print("GPU,&&&&")
 else:
    print("CPU,&&&&")

需要改为:

 device=torch.device('cuda:0')
 if device.type != "cpu":
    print("GPU,&&&&")
 else:
    print("CPU,&&&&")