From ecc77682a056f8dca4f70738965bf4db3b686677 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 30 Jun 2023 11:55:25 -0700 Subject: [PATCH] Check failing UVM test Summary: see if this fixes the error ``` tests/test_uvm_tensor.py::test_uvm_tensor - RuntimeError: CUDA error: invalid device ordinal ``` https://github.com/pytorch/torchsnapshot/actions/runs/5422322107/jobs/9858818161 Differential Revision: D47158196 fbshipit-source-id: 94ca4d13808d1073cac3fba66e57b29e2190bac4 --- tests/test_uvm_tensor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_uvm_tensor.py b/tests/test_uvm_tensor.py index dbd7140..da4053a 100644 --- a/tests/test_uvm_tensor.py +++ b/tests/test_uvm_tensor.py @@ -20,10 +20,12 @@ @pytest.mark.cpu_and_gpu def test_uvm_tensor() -> None: if torch.cuda.is_available() and _UVM_TENSOR_AVAILABLE: + device = torch.device("cuda:0") + torch.cuda.set_device(device) uvm_tensor = torch.rand( (64, 64), out=new_managed_tensor( - torch.empty(0, dtype=torch.float32, device="cuda:0"), + torch.empty(0, dtype=torch.float32, device=device), [64, 64], ), )