CODEWITHSUNDEEP
pythoncupy
Lucas Colley
Lucas Colley

Reputation: 21

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

Upvotes: 1

Views: 413

Answers (1)

Lucas Colley
Lucas Colley

Reputation: 21

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

Upvotes: 1

Related Questions