diff options
| author | Marvin Borner | 2023-09-17 18:15:45 +0200 |
|---|---|---|
| committer | Marvin Borner | 2023-09-17 18:15:45 +0200 |
| commit | 4d42fdd4f9f96018f1fe9024e59278ae78c4f605 (patch) | |
| tree | 394884a56a8ab0dfe96f0dcea910a5d6b396a991 | |
| parent | 3d7dae52244437c1c1acc103f3e124cb2ab1c67e (diff) | |
Model loading fix for DataParallel models
| -rw-r--r-- | swr2_asr/inference.py | 7 |
1 files changed, 6 insertions, 1 deletions
diff --git a/swr2_asr/inference.py b/swr2_asr/inference.py index 3f6a44e..3c58af0 100644 --- a/swr2_asr/inference.py +++ b/swr2_asr/inference.py @@ -66,7 +66,12 @@ def main(config_path: str, file_path: str) -> None: ).to(device) checkpoint = torch.load(inference_config["model_load_path"], map_location=device) - model.load_state_dict(checkpoint["model_state_dict"], strict=True) + + state_dict = { + k[len("module.") :] if k.startswith("module.") else k: v + for k, v in checkpoint["model_state_dict"].items() + } + model.load_state_dict(state_dict, strict=True) model.eval() waveform, sample_rate = torchaudio.load(file_path) # pylint: disable=no-member |