pythoncupy

How to get all available devices for CuPy?

How can I get all available devices for CuPy? I'm looking to write a version of this for CuPy:

if is_torch(xp):
    devices = ['cpu']
    import torch  # type: ignore[import]
    num_cuda = torch.cuda.device_count()
    for i in range(0, num_cuda):
        devices += [f'cuda:{i}']
    if torch.backends.mps.is_available():
        devices += ['mps']
    return devices
Solution

  • I ended up using roughly:

    num_cuda = cupy.cuda.runtime.getDeviceCount()
devices = []
for i in range(0, num_cuda):
    devices += [f'cuda:{i}']
return devices