Fix CI build (#44068)

Summary:
Some of our machines have only 1 device.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/44068

Reviewed By: wanchaol

Differential Revision: D23485730

Pulled By: izdeby

fbshipit-source-id: df6bc0aba18feefc50c56a8f376103352fa2a2ea
This commit is contained in:
iurii zdebskyi
2020-09-02 17:06:50 -07:00
committed by Facebook GitHub Bot
parent 129f406062
commit 2f044d4ee5

View File

@ -156,7 +156,7 @@ class TestForeach(TestCase):
torch._foreach_add_(tensors1, tensors2) torch._foreach_add_(tensors1, tensors2)
# different devices # different devices
if torch.cuda.is_available(): if torch.cuda.is_available() and torch.cuda.device_count() > 1:
tensor1 = torch.zeros(10, 10, device="cuda:0") tensor1 = torch.zeros(10, 10, device="cuda:0")
tensor2 = torch.ones(10, 10, device="cuda:1") tensor2 = torch.ones(10, 10, device="cuda:1")
with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"): with self.assertRaisesRegex(RuntimeError, "Expected all tensors to be on the same device"):