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

Chao1Han · guangyey · pytorchmergebot · commit adb526100a9f · 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_builder.py b/test/distributed/checkpoint/_experimental/test_builder.py
@@ -123,7 +123,7 @@ def test_make_async_checkpointer(self) -> None:
         # Create async checkpointer using factory function with default parameters
         config: CheckpointerConfig = CheckpointerConfig()
         config.staging_config = CheckpointStagerConfig(
-            use_non_blocking_copy=torch.cuda.is_available(),
+            use_non_blocking_copy=torch.accelerator.is_available(),
             use_pinned_memory=torch.cuda.is_available(),
         )
         checkpointer = make_async_checkpointer(config=config, rank_info=self.rank_info)

Original file line number	Diff line number	Diff line change
`@@ -123,7 +123,7 @@ def test_make_async_checkpointer(self) -> None:`
`123`	`123`	`# Create async checkpointer using factory function with default parameters`
`124`	`124`	`config: CheckpointerConfig = CheckpointerConfig()`
`125`	`125`	`config.staging_config = CheckpointStagerConfig(`
`126`		`- use_non_blocking_copy=torch.cuda.is_available(),`
	`126`	`+ use_non_blocking_copy=torch.accelerator.is_available(),`
`127`	`127`	`use_pinned_memory=torch.cuda.is_available(),`
`128`	`128`	`)`
`129`	`129`	`checkpointer = make_async_checkpointer(config=config, rank_info=self.rank_info)`