acc comment

Author: Chao1Han

test/distributed/checkpoint/_experimental/test_builder.py

@@ -124,7 +124,7 @@ def test_make_async_checkpointer(self) -> None:
         config: CheckpointerConfig = CheckpointerConfig()
         config.staging_config = CheckpointStagerConfig(
             use_non_blocking_copy=torch.accelerator.is_available(),
-            use_pinned_memory=torch.cuda.is_available(),
+            use_pinned_memory=torch.accelerator.is_available(),
         )
         checkpointer = make_async_checkpointer(config=config, rank_info=self.rank_info)
 

test/distributed/checkpoint/_experimental/test_staging.py

@@ -108,7 +108,7 @@ def test_different_option_combinations(self) -> None:
             # Only async staging
             test_cases.append(
                 CheckpointStagerConfig(
-                    use_pinned_memory=torch.cuda.is_available(),
+                    use_pinned_memory=torch.accelerator.is_available(),
                     use_shared_memory=False,
                     use_async_staging=True,
                     use_non_blocking_copy=False,
@@ -117,7 +117,7 @@ def test_different_option_combinations(self) -> None:
             # Only CUDA non-blocking copy
             test_cases.append(
                 CheckpointStagerConfig(
-                    use_pinned_memory=torch.cuda.is_available(),
+                    use_pinned_memory=torch.accelerator.is_available(),
                     use_shared_memory=False,
                     use_async_staging=False,
                     use_non_blocking_copy=torch.accelerator.is_available(),
@@ -129,7 +129,7 @@ def test_different_option_combinations(self) -> None:
                 stager = DefaultStager(options)
 
                 # Test staging works with these options
-                if options.use_async_staging and torch.cuda.is_available():
+                if options.use_async_staging and torch.accelerator.is_available():
                     result = stager.stage(self.state_dict)
                     self.assertIsInstance(result, Future)
                     staged_dict = result.result()
@@ -183,7 +183,7 @@ def test_multiple_staging_operations(self) -> None:
         """Test multiple staging operations with the same stager."""
         options = CheckpointStagerConfig(
             use_async_staging=False,
-            use_pinned_memory=torch.cuda.is_available(),
+            use_pinned_memory=torch.accelerator.is_available(),
             use_shared_memory=False,
             use_non_blocking_copy=torch.accelerator.is_available(),
         )