From 185b3c961d72c774ad4030b361c75ebe69b06c31 Mon Sep 17 00:00:00 2001
From: whs <wanghaoshuang@baidu.com>
Date: Mon, 4 Jul 2022 17:38:08 +0800
Subject: [PATCH] Fix saving API when filename is not endswith 'pdmodel'
 (#1239)

---
 paddleslim/auto_compression/compressor.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/paddleslim/auto_compression/compressor.py b/paddleslim/auto_compression/compressor.py
index 7d21bbe7..5885387f 100644
--- a/paddleslim/auto_compression/compressor.py
+++ b/paddleslim/auto_compression/compressor.py
@@ -816,8 +816,14 @@ class AutoCompression:
             for name in test_program_info.feed_target_names
         ]
 
-        model_name = '.'.join(self.model_filename.split(
-            '.')[:-1]) if self.model_filename is not None else 'model'
+        model_name = None
+        if self.model_filename is None:
+            model_name = "model"
+        elif self.model_filename.endswith(".pdmodel"):
+            model_name = self.model_filename.rsplit(".", 1)[0]
+        else:
+            model_name = self.model_filename
+
         path_prefix = os.path.join(model_dir, model_name)
         paddle.static.save_inference_model(
             path_prefix=path_prefix,
-- 
GitLab