python - 如何使用 pytorch 列出所有当前可用的 GPU？

Question

我知道我可以使用 访问当前的 GPU torch.cuda.current_device()，但是如何获取所有当前可用的 GPU 的列表？

标签： pythonpytorchgpu

Answer 1

您可以通过以下方式列出所有可用的 GPU：

>>> import torch
>>> available_gpus = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]
>>> available_gpus
[<torch.cuda.device object at 0x7f2585882b50>]