Update test/distributed/checkpoint/_experimental/test_staging.py

Chao1Han · guangyey · pytorchmergebot · commit 12e06c2ad9d0 · 2025-07-24T02:51:08.000Z
Co-authored-by: Yu, Guangye &lt;106960996+guangyey@users.noreply.github.com&gt;
diff --git a/test/distributed/checkpoint/_experimental/test_staging.py b/test/distributed/checkpoint/_experimental/test_staging.py
@@ -120,7 +120,7 @@ def test_different_option_combinations(self) -> None:
                     use_pinned_memory=torch.cuda.is_available(),
                     use_shared_memory=False,
                     use_async_staging=False,
-                    use_non_blocking_copy=torch.cuda.is_available(),
+                    use_non_blocking_copy=torch.accelerator.is_available(),
                 )
             )
 

Original file line number	Diff line number	Diff line change
`@@ -120,7 +120,7 @@ def test_different_option_combinations(self) -> None:`
`120`	`120`	`use_pinned_memory=torch.cuda.is_available(),`
`121`	`121`	`use_shared_memory=False,`
`122`	`122`	`use_async_staging=False,`
`123`		`- use_non_blocking_copy=torch.cuda.is_available(),`
	`123`	`+ use_non_blocking_copy=torch.accelerator.is_available(),`
`124`	`124`	`)`
`125`	`125`	`)`
`126`	`126`