Add tags for onnx tensor descriptors (#10502)
Summary:
We missed 2 places to add tags when we create tensor descriptors.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/10502
Reviewed By: Maratyszcza
Differential Revision: D9312075
Pulled By: yinghai
fbshipit-source-id: 329e83ec5470b0a778d2eda525dd6f2143facbdf
diff --git a/caffe2/operators/onnxifi_op.cc b/caffe2/operators/onnxifi_op.cc
index 3cdf252..3cfb972 100644
--- a/caffe2/operators/onnxifi_op.cc
+++ b/caffe2/operators/onnxifi_op.cc
@@ -74,6 +74,7 @@
const auto& input_tensor = Input(i);
const auto& tensor_dims = input_tensor.dims();
auto& tensor_descriptor = input_desc_.at(i);
+ tensor_descriptor.tag = ONNXIFI_TAG_TENSOR_DESCRIPTOR_V1;
tensor_descriptor.dataType = ONNXIFI_DATATYPE_FLOAT32;
tensor_descriptor.memoryType = ONNXIFI_MEMORY_TYPE_CPU;
tensor_descriptor.dimensions = tensor_dims.size();
@@ -89,6 +90,7 @@
SetOutputShape(i, &tensor_dims);
output_tensor->Resize(tensor_dims);
auto& tensor_descriptor = output_desc_.at(i);
+ tensor_descriptor.tag = ONNXIFI_TAG_TENSOR_DESCRIPTOR_V1;
tensor_descriptor.dataType = ONNXIFI_DATATYPE_FLOAT32;
tensor_descriptor.memoryType = ONNXIFI_MEMORY_TYPE_CPU;
tensor_descriptor.dimensions = tensor_dims.size();
