From d56759238ba63847375bf3a68cc428caef7c0678 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Fri, 1 Nov 2024 15:43:58 -0300 Subject: [PATCH] Fix test. --- test/test_operations.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/test_operations.py b/test/test_operations.py index cc3a73c4580..892a02ddb0b 100644 --- a/test/test_operations.py +++ b/test/test_operations.py @@ -2993,8 +2993,7 @@ def test_dlpack_xla_to_pytorch_cuda(self): @onlyIfPJRTDeviceIsCUDA def test_dlpack_xla_to_pytorch_cuda_protocol_conversion(self): xla_t1 = torch.arange(5).to(xm.xla_device()) - caps_t1 = torch.utils.dlpack.to_dlpack(xla_t1) - cuda_t1 = torch.utils.dlpack.from_dlpack(caps_t1) + cuda_t1 = torch.utils.dlpack.from_dlpack(xla_t1) self.assertEqual(cuda_t1.device.type, 'cuda') self.assertEqual(cuda_t1.device.index, xla_t1.device.index) cuda_t1[0] = cuda_t1[0] + 20