From b0d68d9b9c819c6df70d694a5d1b4be17f04a449 Mon Sep 17 00:00:00 2001 From: winter-wang <1030748926@qq.com> Date: Tue, 27 Apr 2021 06:52:18 +0000 Subject: [PATCH] update python paddle_trt demo --- python/paddle_trt/infer_trt_ernie.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/paddle_trt/infer_trt_ernie.py b/python/paddle_trt/infer_trt_ernie.py index 3766fd5256a7e..d1e7c7be45ae4 100644 --- a/python/paddle_trt/infer_trt_ernie.py +++ b/python/paddle_trt/infer_trt_ernie.py @@ -9,7 +9,10 @@ # this is a simple resnet block for dynamci test. def init_predictor(args): - config = Config('./model') + if args.model_dir is not "": + config = Config(args.model_dir) + else: + config = Config(args.model_file, args.params_file) config.enable_memory_optim() config.enable_use_gpu(100, 0)