pytorch - 如何使用 CPU 运行 pytorch 项目？

Question

Pytorch 项目应该在 GPU 上运行。我只想在我的笔记本电脑上使用 CPU 运行它。有很多地方调用.cuda()模型、张量等，在 cuda 不可用时无法执行。是否可以在不更改任何地方的代码的情况下做到这一点？

标签： pytorch

Answer 1

这是我能想到的最简单的解决方法：

1. 将以下行放在代码顶部附近：

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

2. 进行全局替换。更改.cuda()为.to(device)，其中device是步骤 1 中设置的变量。