选择服务器上的第二个 GPU

Question

我正在服务器上运行代码。那里有 2 个 GPU，第一个很忙。然而，我找不到在它们之间切换的方法。如果这很重要的话我正在使用 pytorch。应修改以下代码行：

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

修改只能在这里说明。谢谢。

Answer 1

cuda默认情况下选择cuda:0，切换到其他 GPU 可以通过cuda:1

所以，你的行变成：

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

您可以阅读有关CUDA 语义的更多信息。