fix bug

2023-04-17 13:42:50 +08:00 · 2023-04-17 13:42:50 +08:00 · e64ff9facd
parent 7a71b0dd12
commit e64ff9facd
1 changed files with 3 additions and 0 deletions
--- a/finetune.py
+++ b/finetune.py
@ -172,6 +172,9 @@ if not ft_config.skip:
    else:
        trainer.train()

+    # Restore old model state dict
+    model.state_dict = old_state_dict
+
    print('Train completed.')

 # Save Model