diff options
| author |  Jie <jiej@nvidia.com> | 2018-07-11 01:49:16 -0700 |
|---|---|---|
| committer | Jie <jiej@nvidia.com> | 2018-07-11 01:49:16 -0700 |
| commit | 21c4e58510b7126b39d50220716f92ce1716ed2a (patch) | |
| tree | 43367d10308d052d2e782027244059c2debfbb6d /tensorflow/contrib/tensorrt/ops | |
| parent | 6af33fb0436e8c9abad792f33a17e8f662e7311d (diff) | |
enabling shape inference function after cherry picking the fix
Diffstat (limited to 'tensorflow/contrib/tensorrt/ops')
| -rw-r--r-- | tensorflow/contrib/tensorrt/ops/trt_engine_op.cc | 5 |
1 files changed, 2 insertions, 3 deletions
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 |