diff --git a/demo/prune/train.py b/demo/prune/train.py
index 138bf55c5072cdb21b34d3d78947b69593b0f89a..08ec41e0048610a94ec9a7f6eb032fe998c28925 100644
--- a/demo/prune/train.py
+++ b/demo/prune/train.py
@@ -240,8 +240,8 @@ def compress(args):
         if args.save_inference:
             infer_model_path = os.path.join(args.model_path, "infer_models",
                                             str(i))
-            paddle.static.save_inference_model(infer_model_path, ["image"],
-                                               [out], exe, pruned_val_program)
+            paddle.fluid.io.save_inference_model(infer_model_path, ["image"],
+                                                 [out], exe, pruned_val_program)
             _logger.info("Saved inference model into [{}]".format(
                 infer_model_path))