python - 如何从pytorch中的gpu返回cpu？

Question

我在 pytorch 中有一个文本分类器，我想使用 GPU 来提高运行速度。我已经使用这部分代码来检查 CUDA 并使用它：

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    my_rnn_model = nn.DataParallel(my_rnn_model)
if torch.cuda.is_available():
    my_rnn_model.cuda()

现在我想返回使用 cpu（而不是 gpu）。所以我清除了这部分代码。但它不起作用，我收到此错误：

RuntimeError: cuda runtime error (8) : invalid device function at /opt/conda/conda-bld/pytorch_1503963423183/work/torch/lib/THC/THCTensorCopy.cu:204

请您指导我如何才能返回到 cpu 运行？

标签： pythongpucpupytorch

Answer 1

您可以设置要使用的 GPU 设备：

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

在您的情况下，您可以使用以下命令返回 CPU：

torch.device('cpu')