Преглед изворни кода

call enable_use_gpu() first to use TensorRT engine

gaotingquan пре 8 месеци
родитељ
комит
696089a110
1 измењених фајлова са 2 додато и 1 уклоњено
  1. 2 1
      paddlex/inference/models/common/static_infer.py

+ 2 - 1
paddlex/inference/models/common/static_infer.py

@@ -484,8 +484,9 @@ class StaticInfer(object):
             config = lazy_paddle.inference.Config(str(model_file), str(params_file))
         else:
             config = lazy_paddle.inference.Config(str(model_file), str(params_file))
-
             config.set_optim_cache_dir(str(cache_dir / "optim_cache"))
+            # call enable_use_gpu() first to use TensorRT engine
+            config.enable_use_gpu(100, self._option.device_id)
             for func_name in self._option.trt_cfg_setting:
                 assert hasattr(
                     config, func_name