Skip to content

Commit 79a0777

Browse files
authored
Merge pull request PaddlePaddle#124 from winter-wang/master
update python paddle_trt demo
2 parents 1054925 + b0d68d9 commit 79a0777

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

python/paddle_trt/infer_trt_ernie.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,10 @@
99

1010
# this is a simple resnet block for dynamci test.
1111
def init_predictor(args):
12-
config = Config('./model')
12+
if args.model_dir is not "":
13+
config = Config(args.model_dir)
14+
else:
15+
config = Config(args.model_file, args.params_file)
1316
config.enable_memory_optim()
1417
config.enable_use_gpu(100, 0)
1518

0 commit comments

Comments
 (0)