kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-05 02:31:33 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[https://nvbugs/5816267][fix] Remove weight tensor holder to release memory earlier (#10876)

Browse Source 
Signed-off-by: Dongxu Yang <78518666+dongxuy04@users.noreply.github.com>
This commit is contained in:
dongxuy04 2026-01-21 16:42:52 -08:00 committed by GitHub
parent 5450485bec
commit 635cbf01ba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorrt_llm/_torch/pyexecutor/model_loader.py
View File

@ -256,6 +256,9 @@ class ModelLoader:
f"Fallback to regular model init: {traceback.format_exc(limit=10)}\n"
)
model = AutoModelForCausalLM.from_config(config)
finally:
if 'memo' in locals():
del memo
model.to("cuda")
rank_model_storage = get_rank_model_storage(model)