if torch.cuda.is_available():
torch.set_default_tensor_type(torch. cuda. FloatTensor)
print ("using cuda:",torch.cuda.get_device_name(0))
pass
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device
检查的语句是torch.cuda.is_available()。如果设备可用,我们将默认的数据类型设为torch.cuda.FloatTensor。然后,我们打印出CUDA找到的设备名。这段标准代码同时设置一个PyTorch设备,在没有默认使用GPU的情况下,允许我们通过代码将数据传入GPU。如果CUDA不可用,则设备会指向CPU。