From 21c4e58510b7126b39d50220716f92ce1716ed2a Mon Sep 17 00:00:00 2001 From: Jie Date: Wed, 11 Jul 2018 01:49:16 -0700 Subject: enabling shape inference function after cherry picking the fix --- tensorflow/contrib/tensorrt/ops/trt_engine_op.cc | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) (limited to 'tensorflow/contrib/tensorrt/ops') diff --git a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc b/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc index 7eaa080586..383635f428 100644 --- a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc +++ b/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc @@ -42,9 +42,8 @@ REGISTER_OP("TRTEngineOp") .Attr("precision_mode: {'FP32', 'FP16', 'INT8', 'INT8CALIB'}") .Attr("calibration_data: string = ''") .Input("in_tensor: InT") - .Output("out_tensor: OutT"); - // TODO(Sami): shape inference not working for concrete input shape - //.SetShapeFn(shape_inference::TRTEngineOpShapeInference); + .Output("out_tensor: OutT") + .SetShapeFn(shape_inference::TRTEngineOpShapeInference); } // namespace tensorflow -- cgit v1.2.3