From e64ff9facdea8856b6859bf388b7d0d6c9366e0f Mon Sep 17 00:00:00 2001
From: John Smith <yfshi123@163.com>
Date: Mon, 17 Apr 2023 13:42:50 +0800
Subject: [PATCH] fix bug

---
 finetune.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/finetune.py b/finetune.py
index 9874d1b..b53be15 100644
--- 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