diff --git a/finetuning.py b/finetuning.py
index 68cfd84..c0800d8 100644
--- a/finetuning.py
+++ b/finetuning.py
@@ -1,5 +1,5 @@
 import torch
-from pocket_flow.gdbp_model import PocketFlowWithEdgeNew, reset_parameters, freeze_parameters
+from pocket_flow.gdbp_model import PocketFlow, reset_parameters, freeze_parameters
 from pocket_flow.utils import Experiment, LoadDataset
 from pocket_flow.utils.transform import *
 #from utils.ParseFile import Protein, parse_sdf_to_dict
@@ -54,4 +54,4 @@ exp.fit_step(
     with_tb=True, logdir='./finetuning_log', schedule_key='loss', num_workers=8, 
     pin_memory=False, follow_batch=[], exclude_keys=[], collate_fn=None, 
     max_edge_num_in_batch=400000
-    )
\ No newline at end of file
+    )