From ae49b2a8e1bc676e6014d8a12ebeec947b655e3e Mon Sep 17 00:00:00 2001
From: 莫拉古 <61447879+yechaoying@users.noreply.github.com>
Date: 星期五, 29 十一月 2024 09:55:43 +0800
Subject: [PATCH] 变量名写错了 (#2249)
---
funasr/train_utils/load_pretrained_model.py | 7 ++++---
1 files changed, 4 insertions(+), 3 deletions(-)
diff --git a/funasr/train_utils/load_pretrained_model.py b/funasr/train_utils/load_pretrained_model.py
index 8ed613c..da2eed5 100644
--- a/funasr/train_utils/load_pretrained_model.py
+++ b/funasr/train_utils/load_pretrained_model.py
@@ -8,6 +8,7 @@
import torch.nn
import torch.optim
import pdb
+import copy
def load_pretrained_model(
@@ -35,11 +36,12 @@
logging.info(f"ckpt: {path}")
if oss_bucket is None:
- src_state = torch.load(path, map_location=map_location)
+ ori_state = torch.load(path, map_location=map_location)
else:
buffer = BytesIO(oss_bucket.get_object(path).read())
- src_state = torch.load(buffer, map_location=map_location)
+ ori_state = torch.load(buffer, map_location=map_location)
+ src_state = copy.deepcopy(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
@@ -94,7 +96,6 @@
)
else:
dst_state[k] = src_state[k_src]
-
else:
print(f"Warning, miss key in ckpt: {k}, {path}")
--
Gitblit v1.9.1