Skip to content

Commit cf84dac

Browse files
nikhilaravifacebook-github-bot
authored andcommitted
fix get cuda device test error
Summary: Cuda test failing on circle with the error `random_ expects 'from' to be less than 'to', but got from=0 >= to=0` This is because the `high` value in `torch.randint` is 1 more than the highest value in the distribution from which a value is drawn. So if there is only 1 cuda device available then the low and high are 0. Reviewed By: gkioxari Differential Revision: D21236669 fbshipit-source-id: 46c312d431c474f1f2c50747b1d5e7afbd7df3a9
1 parent f8acecb commit cf84dac

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

tests/common_testing.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,10 @@ def get_random_cuda_device() -> str:
2828
any device without having to set the device explicitly.
2929
"""
3030
num_devices = torch.cuda.device_count()
31-
rand_device_id = torch.randint(high=num_devices, size=(1,)).item()
32-
return "cuda:%d" % rand_device_id
31+
device_id = (
32+
torch.randint(high=num_devices, size=(1,)).item() if num_devices > 1 else 0
33+
)
34+
return "cuda:%d" % device_id
3335

3436

3537
class TestCaseMixin(unittest.TestCase):

0 commit comments

Comments
 (0)