Merge pull request PaddlePaddle#124 from winter-wang/master

update python paddle_trt demo
zhoutianzi666 · Apr 27, 2021 · 79a0777 · 79a0777
2 parents 1054925 + b0d68d9
commit 79a0777
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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)