如何将 torch.device('cuda' if torch.cuda.is_available() else 'cpu') 编写为完整的 if else 语句？

Question

我是 Pytorch 的初学者，想将这个语句作为一个整体输入 if else 语句：-

torch.device('cuda' if torch.cuda.is_available() else 'cpu')

有人可以帮助我吗？

Answer 1

下面是整个 if-else 语句的代码：

torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if torch.cuda.is_available():
    torch.device('cuda')
else:
    torch.device('cpu')

由于您可能希望稍后存储设备，您可能需要这样的东西：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

这是一篇关于 Python 中三元运算符的帖子和讨论：https : //stackoverflow.com/a/2802748/13985765

从那个帖子：

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')