[ckpt] print more info for debug when loading state_dict (#282)

wenet-e2e · Mar 4, 2024 · 94e91ee · 94e91ee
1 parent 31921e9
commit 94e91ee
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/wespeaker/utils/checkpoint.py b/wespeaker/utils/checkpoint.py
@@ -14,11 +14,17 @@
 # limitations under the License.
 
 import torch
+import logging
 
 
 def load_checkpoint(model: torch.nn.Module, path: str):
     checkpoint = torch.load(path, map_location='cpu')
-    model.load_state_dict(checkpoint, strict=False)
+    missing_keys, unexpected_keys = model.load_state_dict(checkpoint,
+                                                          strict=False)
+    for key in missing_keys:
+        logging.warning('missing tensor: {}'.format(key))
+    for key in unexpected_keys:
+        logging.warning('unexpected tensor: {}'.format(key))
 
 
 def save_checkpoint(model: torch.nn.Module, path: str):