modelscope/FunASR
1
0
Fork 0
mirror of https://github.com/modelscope/FunASR.git synced 2026-10-12 02:19:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: avoid deep-copying the checkpoint in load_pretrained_model (#3729)

Browse Source
This commit is contained in:
YYDH
2026-09-26 18:56:23 +08:00
committed by GitHub
parent 41778c45de
commit 9e7f6af8c3
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
funasr/train_utils/load_pretrained_model.py
+4 -2
View File
@@ -8,7 +8,6 @@ import torch
import torch.nn
import torch.optim
import pdb
import copy
def load_pretrained_model(
@@ -41,7 +40,10 @@ def load_pretrained_model(
buffer = BytesIO(oss_bucket.get_object(path).read())
ori_state = torch.load(buffer, map_location=map_location)
src_state = copy.deepcopy(ori_state)
# ori_state is created by torch.load inside this function, is not shared
# with the caller, and is not read again below; src_state is only read.
# Deep-copying every tensor here doubles peak load memory with no effect.
src_state = ori_state
src_state = src_state["state_dict"] if "state_dict" in src_state else src_state
src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state
src_state = src_state["model"] if "model" in src_state else src_state