diff options
| author | Pherkel | 2023-09-11 22:58:19 +0200 |
|---|---|---|
| committer | Pherkel | 2023-09-11 22:58:19 +0200 |
| commit | 6f5513140f153206cfa91df3077e67ce58043d35 (patch) | |
| tree | 71cee784719a1f9c912a1038824eb6bd26195408 /swr2_asr/train.py | |
| parent | 64dbb9d32a51b1bce6c9de67069dc8f5943a5399 (diff) | |
model loading is broken :(
Diffstat (limited to 'swr2_asr/train.py')
| -rw-r--r-- | swr2_asr/train.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/swr2_asr/train.py b/swr2_asr/train.py index ca70d21..ec25918 100644 --- a/swr2_asr/train.py +++ b/swr2_asr/train.py @@ -263,7 +263,7 @@ def main(config_path: str): prev_epoch = 0 if checkpoints_config["model_load_path"] is not None: - checkpoint = torch.load(checkpoints_config["model_load_path"]) + checkpoint = torch.load(checkpoints_config["model_load_path"], map_location=device) model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) prev_epoch = checkpoint["epoch"] |