python - 如何将 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')

有人可以帮助我吗？

标签： pythonif-statementpytorchgpu

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

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

---

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

从那个帖子：

value_when_true if condition else value_when_false