diff --git a/tests/test_tipc/ernie_text_matching/predict.py b/tests/test_tipc/ernie_text_matching/predict.py index 2621a7126e1c..a93a471bea66 100644 --- a/tests/test_tipc/ernie_text_matching/predict.py +++ b/tests/test_tipc/ernie_text_matching/predict.py @@ -22,6 +22,10 @@ from paddlenlp.data import Pad, Tuple from paddlenlp.datasets import load_dataset from paddlenlp.transformers import AutoTokenizer +from paddlenlp.utils.env import ( + PADDLE_INFERENCE_MODEL_SUFFIX, + PADDLE_INFERENCE_WEIGHTS_SUFFIX, +) from paddlenlp.utils.log import logger @@ -59,8 +63,8 @@ def __init__( self.batch_size = batch_size self.benchmark = benchmark - model_file = os.path.join(model_dir, "inference.pdmodel") - params_file = os.path.join(model_dir, "inference.pdiparams") + model_file = os.path.join(model_dir, f"inference{PADDLE_INFERENCE_MODEL_SUFFIX}") + params_file = os.path.join(model_dir, f"inference{PADDLE_INFERENCE_WEIGHTS_SUFFIX}") if not os.path.exists(model_file): raise ValueError("not find model file path {}".format(model_file)) if not os.path.exists(params_file):