[Hardware] Replace memory related torch.cuda APIs (#37031)

Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
2026-03-16 18:24:48 +08:00
parent 122f75d939
commit 747b068136
8 changed files with 20 additions and 17 deletions
--- a/tests/utils_/test_mem_utils.py
+++ b/tests/utils_/test_mem_utils.py
@@ -29,7 +29,7 @@ def test_memory_profiling():
    def measure_current_non_torch():
        free, total = torch.cuda.mem_get_info()
        current_used = total - free
-        current_torch = torch.cuda.memory_reserved()
+        current_torch = torch.accelerator.memory_reserved()
        current_non_torch = current_used - current_torch
        return current_non_torch