Skip to content

Commit 0c0dc9d

Browse files
committed
Add missed variable
Signed-off-by: Hui Gao <huig@nvidia.com>
1 parent 01530c6 commit 0c0dc9d

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -367,7 +367,8 @@ class DeepGemmFusedMoE(CutlassFusedMoE):
367367
"""
368368

369369
# To reuse pytorch memory segments allocated during graph capture.
370-
allocated_buffer_recorder: dict[str, list[torch.Tensor]] = {}
370+
allocated_buffer_in_graph_pool: dict[str, list[torch.Tensor]] = {}
371+
allocated_buffer_in_runtime: dict[str, torch.Tensor] = {}
371372

372373
def __init__(
373374
self,

0 commit comments

Comments
 (0)