Searched refs:trt_axis (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/ |
| D | convert_nodes.cc | 530 bool use_implicit_batch, int* trt_axis) { in ConvertAxis() argument
547 *trt_axis = use_implicit_batch ? tf_axis - 1 : tf_axis; in ConvertAxis()
2656 int trt_axis; in ConvertExpandDims() local
2658 params->use_implicit_batch, &trt_axis)); in ConvertExpandDims()
2664 input_tensor.tensor(), dims, trt_axis, params, &output_tensor)); in ConvertExpandDims()
2667 input_dims.insert(input_dims.begin() + trt_axis, 1); in ConvertExpandDims()
2819 int trt_axis; in ConvertSqueeze() local
2821 params->use_implicit_batch, &trt_axis)); in ConvertSqueeze()
2823 if (input_dims[trt_axis] != -1 && input_dims[trt_axis] != 1) { in ConvertSqueeze()
2825 "Dimension ", tf_axis, " with size ", input_dims[trt_axis], in ConvertSqueeze()
[all …]
|