Browse Source

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

GitOrigin-RevId: b830706a89
tags/v1.8.0
Megvii Engine Team 3 years ago
parent
commit
f4f20046c4
1 changed files with 1 additions and 0 deletions
  1. +1
    -0
      src/tensorrt/impl/tensorrt_runtime_opr.cpp

+ 1
- 0
src/tensorrt/impl/tensorrt_runtime_opr.cpp View File

@@ -107,6 +107,7 @@ TensorRTRuntimeOpr::TensorRTRuntimeOpr(
void TensorRTRuntimeOpr::get_output_var_shape(
const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const {
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());
auto get_mgb_shape = [&](int binding_idx) -> TensorShape {
auto dims = m_engine->getBindingDimensions(binding_idx);


Loading…
Cancel
Save