Skip to content

Commit dce6e64

Browse files
authored
[Fix] Fix the warning method that calls the trt logger (#1850)
* Fix the warning method that calls the trt logger * Fix the warning method that calls the trt logger
1 parent 50bd7e9 commit dce6e64

File tree

1 file changed

+8
-3
lines changed

1 file changed

+8
-3
lines changed

mmdeploy/backend/tensorrt/utils.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -152,10 +152,15 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto],
152152
else:
153153
os.environ.pop('CUDA_DEVICE')
154154

155+
# build a mmdeploy logger
156+
logger = get_root_logger()
155157
load_tensorrt_plugin()
158+
159+
# build a tensorrt logger
160+
trt_logger = trt.Logger(log_level)
161+
156162
# create builder and network
157-
logger = trt.Logger(log_level)
158-
builder = trt.Builder(logger)
163+
builder = trt.Builder(trt_logger)
159164

160165
# TODO: use TorchAllocator as builder.gpu_allocator
161166

@@ -164,7 +169,7 @@ def from_onnx(onnx_model: Union[str, onnx.ModelProto],
164169
network = builder.create_network(EXPLICIT_BATCH)
165170

166171
# parse onnx
167-
parser = trt.OnnxParser(network, logger)
172+
parser = trt.OnnxParser(network, trt_logger)
168173

169174
if isinstance(onnx_model, str):
170175
parse_valid = parser.parse_from_file(onnx_model)

0 commit comments

Comments
 (0)