diff --git a/slim/prune/train_prune.py b/slim/prune/train_prune.py
index 364130f7ad2ea2e2ef733cd6391deb8e77fdf893..06e1658f1a3f721842fbe780820103aceac87a16 100644
--- a/slim/prune/train_prune.py
+++ b/slim/prune/train_prune.py
@@ -47,8 +47,7 @@ from eval_prune import evaluate
 from vis import visualize
 from utils import dist_utils
 
-from paddleslim.prune import Pruner
-from paddleslim.prune.io import *
+from paddleslim.prune import Pruner, save_model
 from paddleslim.analysis import flops
 
 def parse_args():