Parcourir la source

call enable_use_gpu() first to use TensorRT engine

gaotingquan il y a 8 mois
Parent
commit
696089a110
1 fichiers modifiés avec 2 ajouts et 1 suppressions
  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