diff --git a/onnxruntime/core/providers/cpu/tensor/split.cc b/onnxruntime/core/providers/cpu/tensor/split.cc index db9315356a..2af92504ec 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.cc +++ b/onnxruntime/core/providers/cpu/tensor/split.cc @@ -18,6 +18,7 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( std::vector{ DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), Split); @@ -29,6 +30,7 @@ ONNX_CPU_OPERATOR_KERNEL( std::vector{ DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}), Split); @@ -79,6 +81,8 @@ Status Split::Compute(OpKernelContext* context) const { status = ComputeImpl(*context, input); else if (input.IsDataType()) status = ComputeImpl(*context, input); + else if (input.IsDataType()) + status = ComputeImpl(*context, input); else if (input.IsDataTypeString()) status = ComputeImpl(*context, input); else @@ -92,7 +96,7 @@ inline void copy_data(const T* src, T* dst, size_t count) { memcpy(dst, src, count * sizeof(T)); } -template<> +template <> inline void copy_data(const std::string* src, std::string* dst, size_t count) { const std::string* end = src + count; std::copy(src, end, dst); diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc index 1957f40908..1ff8c4d79d 100644 --- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc @@ -11,7 +11,6 @@ template using ShapeAndData = std::pair, const std::vector>; using ShapeAndFloatData = ShapeAndData; -using ShapeAndInt32Data = ShapeAndData; using ShapeAndStringData = ShapeAndData; using ExpectResult = OpTester::ExpectResult; @@ -65,16 +64,17 @@ TEST(SplitOperatorTest, Axis0EqualSplitFloat) { RunTest(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM } -TEST(SplitOperatorTest, Axis0EqualSplitInt32) { +template ::value, T>::type> +static void SplitTestInt() { const int64_t axis = 0; - std::vector outputs; + std::vector> outputs; // input shape and data - ShapeAndInt32Data input = {{4, 2}, // shape - {1, 2, - 3, 4, - 5, 6, - 7, 8}}; + ShapeAndData input = {{4, 2}, // shape + {1, 2, + 3, 4, + 5, 6, + 7, 8}}; outputs.push_back({{2, 2}, {1, 2, @@ -84,7 +84,15 @@ TEST(SplitOperatorTest, Axis0EqualSplitInt32) { {5, 6, 7, 8}}); - RunTest(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, {}, input, outputs, false); //TensorRT parser: Assertion failed: axis != BATCH_DIM +} + +TEST(SplitOperatorTest, Axis0EqualSplitInt32) { + SplitTestInt(); +} + +TEST(SplitOperatorTest, Axis0EqualSplitInt64) { + SplitTestInt(); } TEST(SplitOperatorTest, Axis0EqualSplitString) {