frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[dynamo][easy] Support torch.accelerator.current_accelerator (#165734)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165734
Approved by: https://github.com/Skylion007
This commit is contained in:
Animesh Jain
2025-10-17 09:46:53 -07:00
committed by PyTorch MergeBot
parent 86ebce1766
commit c18ddfc572
2 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
test/dynamo/test_repros.py
View File

@ -8101,6 +8101,14 @@ class ReproTestsDevice(torch._dynamo.test_case.TestCase):
res = gm(x, y) res = gm(x, y)
self.assertEqual(res, ref) self.assertEqual(res, ref)
def test_current_accelerator(self):
@torch.compile(backend="eager", fullgraph=True)
def fn(x):
torch.accelerator.current_accelerator()
return x + 1
self.assertEqual(fn(torch.ones(3)), torch.ones(3) + 1)
instantiate_parametrized_tests(ReproTests) instantiate_parametrized_tests(ReproTests)

1
torch/_dynamo/variables/torch.py
View File

@ -146,6 +146,7 @@ REWRITE_OPS_TO_TENSOR_SIZE_METHOD = dict.fromkeys(
constant_fold_functions_need_guards = [ constant_fold_functions_need_guards = [
torch.accelerator.current_device_index, torch.accelerator.current_device_index,
torch.accelerator.current_accelerator,
torch.cuda.current_device, torch.cuda.current_device,
torch.cuda.is_initialized, torch.cuda.is_initialized,
torch.xpu.current_device, torch.xpu.current_device,