[TEST][Dynamo] fix test_deviceguard.py (#126240)

The `test_device_guard.py` was improperly set up, so there were failures on multi-GPU machines. By design the `DeviceGuard` should keep `idx` the same even after it was applied.

Pull Request resolved: #126240
Approved by: https://github.com/jansel

test/dynamo/test_deviceguard.py

@@ -20,7 +20,7 @@ def setUp(self):
         self.device_interface = Mock()
 
         self.device_interface.exchange_device = Mock(return_value=0)
-        self.device_interface.maybe_exchange_device = Mock(return_value=0)
+        self.device_interface.maybe_exchange_device = Mock(return_value=1)
 
     def test_device_guard(self):
         device_guard = DeviceGuard(self.device_interface, 1)
@@ -32,7 +32,7 @@ def test_device_guard(self):
 
         self.device_interface.maybe_exchange_device.assert_called_once_with(0)
         self.assertEqual(device_guard.prev_idx, 0)
-        self.assertEqual(device_guard.idx, 0)
+        self.assertEqual(device_guard.idx, 1)
 
     def test_device_guard_no_index(self):
         device_guard = DeviceGuard(self.device_interface, None)
@@ -70,7 +70,7 @@ def test_device_guard(self):
 
         self.assertEqual(torch.cuda.current_device(), current_device)
         self.assertEqual(device_guard.prev_idx, 0)
-        self.assertEqual(device_guard.idx, 0)
+        self.assertEqual(device_guard.idx, 1)
 
     def test_device_guard_no_index(self):
         current_device = torch.cuda.current_device()