From cf84dacf2ed3d38b5d79070c23993a5265a4ee4b Mon Sep 17 00:00:00 2001 From: Nikhila Ravi Date: Fri, 24 Apr 2020 16:10:11 -0700 Subject: [PATCH] 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 --- tests/common_testing.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/common_testing.py b/tests/common_testing.py index 7f6898a8..ddf1fc3d 100644 --- a/tests/common_testing.py +++ b/tests/common_testing.py @@ -28,8 +28,10 @@ def get_random_cuda_device() -> str: any device without having to set the device explicitly. """ num_devices = torch.cuda.device_count() - rand_device_id = torch.randint(high=num_devices, size=(1,)).item() - return "cuda:%d" % rand_device_id + device_id = ( + torch.randint(high=num_devices, size=(1,)).item() if num_devices > 1 else 0 + ) + return "cuda:%d" % device_id class TestCaseMixin(unittest.TestCase):