为什么torch.cuda.device的设备参数可能大于torch.cuda.device_count（）？

Question

是代码：

torch.cuda.device_count()
# output: 1
d2 = torch.cuda.device(2)
# d2.idx = 2

是在其他设备上运行的模型吗？

Answer 1

torch.cuda.device_count()返回可用GPU的数量，而torch.cuda.device(device)是context-manager。它有助于正确处理资源，即在使用后自动设置和释放资源。因此，您在此函数中传递的参数并不表示您正在切换/设置该设备。为了设置当前设备，您应该使用torch.cuda.set_device(device)。虽然，根据官方文档here，

  

不建议使用此功能，而建议使用设备。多数情况   情况下，最好使用CUDA_VISIBLE_DEVICES环境变量。

要详细了解torch.cuda.device（）及其工作原理，请进行this官方pytorch讨论。