[CPU] Use accelerator memory API in DiffusionGemma (#58964)
Signed-off-by: zhejiangxiaomai <zhenhui.zhao@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
zhao, zhenhui committed
Sep 28, 2026 at 13:31 UTC
1668e5976db7457631b84a6a570c64319581b1e7
1 file changed
+1
-1
vllm/model_executor/models/diffusion_gemma.py
+1
-1
@@ -1493,7 +1493,7 @@ class DiffusionSampler:
1493
# [tile * W, vocab] copies, so a tile is also bounded by free memory.
1494
widths_np = states.canvas_width_np[decode_slots_np]
1495
order = np.argsort(widths_np, kind="stable")
1496
- free = current_platform.mem_get_info()[0] if num_decode > 0 else 0
1496
+ free = torch.accelerator.get_memory_info()[0] if num_decode > 0 else 0
1497
run_start = 0
1498
while run_start < num_decode:
1499
W = int(widths_np[order[run_start]])