Fix on XPU

guangyey · guangyey · commit 3cc73a6bfc81 · 2025-05-05T09:59:08.000Z
ghstack-source-id: 9626b07 Pull Request resolved: #152812
diff --git a/test/test_xpu.py b/test/test_xpu.py
@@ -334,13 +334,13 @@ def test_device_context_manager(self):
         self.assertEqual(torch.xpu.current_device(), prev_device)
 
     @unittest.skipIf(not TEST_MULTIXPU, "only one GPU detected")
-    def test_device_context_manager_with_set_device(self):
+    def test_multi_device_context_manager(self):
         src_device = 0
         dst_device = 1
         torch.xpu.set_device(src_device)
         with torch.accelerator.device_index(dst_device):
             self.assertEqual(torch.xpu.current_device(), 1)
-        self.assertEqual(torch.xpu.set_device(), src_device)
+        self.assertEqual(torch.xpu.current_device(), src_device)
 
     def test_stream_context_manager(self):
         prev_stream = torch.xpu.current_stream()
