Skip to content

Commit 7314cf4

Browse files
dvrogozhpytorchmergebot
authored andcommitted
torch/accelerator: fix device type comparison (#143541)
This was failing without the fix: ``` python -c 'import torch; d=torch.device("xpu:0"); torch.accelerator.current_stream(d)' ``` with: ``` ValueError: xpu doesn't match the current accelerator xpu. ``` CC: @guangyey, @EikanWang Pull Request resolved: #143541 Approved by: https://github.com/guangyey, https://github.com/albanD
1 parent 434e0c2 commit 7314cf4

File tree

2 files changed

+12
-1
lines changed

2 files changed

+12
-1
lines changed

test/test_accelerator.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,17 @@ def test_generic_stream_behavior(self):
6868
self.assertTrue(event.query())
6969
self.assertEqual(c_acc.cpu(), c)
7070

71+
def test_current_stream_query(self):
72+
s = torch.accelerator.current_stream()
73+
self.assertEqual(torch.accelerator.current_stream(s.device), s)
74+
self.assertEqual(torch.accelerator.current_stream(s.device.index), s)
75+
self.assertEqual(torch.accelerator.current_stream(str(s.device)), s)
76+
other_device = torch.device("cpu")
77+
with self.assertRaisesRegex(
78+
ValueError, "doesn't match the current accelerator"
79+
):
80+
torch.accelerator.current_stream(other_device)
81+
7182

7283
if __name__ == "__main__":
7384
run_tests()

torch/accelerator/_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ def _get_device_index(device: _device_t, optional: bool = False) -> int:
1111
device = torch.device(device)
1212
device_index: Optional[int] = None
1313
if isinstance(device, torch.device):
14-
if torch.accelerator.current_accelerator() != device.type:
14+
if torch.accelerator.current_accelerator().type != device.type:
1515
raise ValueError(
1616
f"{device.type} doesn't match the current accelerator {torch.accelerator.current_accelerator()}."
1717
)

0 commit comments

Comments
 (0)