add test/prototype/test_parq.py

zxd1997066 · zxd1997066 · commit 346a33962586 · 2025-12-02T15:35:32.000+08:00
diff --git a/test/prototype/test_parq.py b/test/prototype/test_parq.py
@@ -51,7 +51,7 @@
     torch_version_at_least,
 )
 
-_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+_DEVICE = torch.device(torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu")
 
 
 class M(nn.Module):

Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@`
`51`	`51`	`torch_version_at_least,`
`52`	`52`	`)`
`53`	`53`
`54`		`-_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")`
	`54`	`+_DEVICE = torch.device(torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu")`
`55`	`55`
`56`	`56`
`57`	`57`	`class M(nn.Module):`