提交 f4f20046 编写于 作者: M Megvii Engine Team

fix(mgb): fix tensorrt runtimeopr get output var shape bug

GitOrigin-RevId: b830706a891e4859d8e6a0d3db337a2941577168
上级 82be0aac
...@@ -107,6 +107,7 @@ TensorRTRuntimeOpr::TensorRTRuntimeOpr( ...@@ -107,6 +107,7 @@ TensorRTRuntimeOpr::TensorRTRuntimeOpr(
void TensorRTRuntimeOpr::get_output_var_shape( void TensorRTRuntimeOpr::get_output_var_shape(
const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const { const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const {
auto batch = inp_shape.at(0)[0]; auto batch = inp_shape.at(0)[0];
m_manager.clear_trt_context();
m_manager.create_trt_context(this->comp_node(), inp_shape, m_engine.get()); m_manager.create_trt_context(this->comp_node(), inp_shape, m_engine.get());
auto get_mgb_shape = [&](int binding_idx) -> TensorShape { auto get_mgb_shape = [&](int binding_idx) -> TensorShape {
auto dims = m_engine->getBindingDimensions(binding_idx); auto dims = m_engine->getBindingDimensions(binding_idx);
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册