FillVectorsFromInput Starts must be a 1-D array

Description

Let me just preface that this is not a ONNX Surgeon bug, but probably something on my part that I don't know how to solve. Also, since I've been searching the internet for the past few days, this seems like the best place to ask the question. I'm trying to remove ArrayFeatureExtractor layer from my network and replace it with the combination of Split and Slice layers. ONNX Surgeon part goes without errors and I get my new network with input "starts" and "ends" being a 1-D array for the Slice layer. Later, when I run it through onnx runtime API, I get the following output:

2022-01-27 13:44:57.527266048 [W:onnxruntime:, graph.cc:122 MergeShapeInfo] Error merging shape info for output. 'split_out_1' source:{,1} target:{1}. Falling back to lenient merge.
2022-01-27 13:44:57.527293714 [W:onnxruntime:, graph.cc:122 MergeShapeInfo] Error merging shape info for output. 'split_out_2' source:{,1} target:{1}. Falling back to lenient merge.
2022-01-27 13:44:57.527305561 [W:onnxruntime:, graph.cc:122 MergeShapeInfo] Error merging shape info for output. 'split_out_3' source:{,1} target:{1}. Falling back to lenient merge.
2022-01-27 13:44:57.536553595 [E:onnxruntime:, sequential_executor.cc:346 Execute] Non-zero status code returned while running Slice node. Name:'' Status Message: slice.cc:153 FillVectorsFromInput Starts must be a 1-D array
Traceback (most recent call last):
File "inference.py", line 13, in
outputs = ort_sess.run(None, {'float_input': x})
File "/home/alex/anaconda3/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 192, in run
return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Slice node. Name:'' Status Message: slice.cc:153 FillVectorsFromInput Starts must be a 1-D array

I'm posting both of my networks, my python code for onnx surgeon and my python code for simple inference.

import onnx_graphsurgeon as gs
import onnx
import numpy as np

graph = gs.import_onnx(onnx.load("knn-fixed.onnx"))

inputs_array = []

for node in graph.nodes:
    if node.op == "ArrayFeatureExtractor":
        inputs_array.append(node.inputs)

class_tensor = inputs_array[0][0]

flatten_node = [node for node in graph.nodes if node.op == "Flatten"][0]

split_shape = np.array([0]).shape

split_out_1 = gs.Variable("split_out_1", dtype = np.int64, shape = split_shape)
split_out_2 = gs.Variable("split_out_2", dtype = np.int64, shape = split_shape)
split_out_3 = gs.Variable("split_out_3", dtype = np.int64, shape = split_shape)
split_node = gs.Node(op = "Split", inputs = flatten_node.outputs, outputs = [split_out_1, split_out_2, split_out_3])
split_node.attrs["axis"] = 1
split_node.attrs["split"] = [1, 1, 1]

slice_out_1 = gs.Variable("slice_out_1", dtype = np.int64)
slice_out_2 = gs.Variable("slice_out_2", dtype = np.int64)
slice_out_3 = gs.Variable("slice_out_3", dtype = np.int64)

slice_node_1 = gs.Node(op = "Slice", inputs = [class_tensor, split_out_1, split_out_1], outputs = [slice_out_1])
slice_node_2 = gs.Node(op = "Slice", inputs = [class_tensor, split_out_2, split_out_2], outputs = [slice_out_2])
slice_node_3 = gs.Node(op = "Slice", inputs = [class_tensor, split_out_3, split_out_3], outputs = [slice_out_3])

concat_out = gs.Variable("concat_out", dtype = np.int64)
concat_node = gs.Node(op = "Concat", inputs = [slice_out_1, slice_out_2, slice_out_3], outputs = [concat_out])
concat_node.attrs["axis"] = 0

graph.outputs = concat_node.outputs

graph.nodes.append(split_node)

graph.nodes.append(slice_node_1)
graph.nodes.append(slice_node_2)
graph.nodes.append(slice_node_3)

graph.nodes.append(concat_node)

graph.cleanup().toposort()

onnx.save(gs.export_onnx(graph), "knn-fixed-2.onnx")

Code for simple inference:

import onnx
import pandas as pd
import numpy as np

df = pd.read_csv("../some_path/keypoints.csv")
test = df.iloc[0].drop(['item', 'class'])
test = np.array(test).astype(np.float32).reshape(-1, 266)

import onnxruntime as ort
x = test
y = df.iloc[0]['class']
ort_sess = ort.InferenceSession('knn-fixed-2.onnx')
outputs = ort_sess.run(None, {'float_input': x})

print(f"Predicted: {outputs}")

Again, I understand how this may not be an appropriate place to ask the question, but I couldn't find any forum dedicated to this and stackoverflow doesn't seem to be the best place when it comes to ONNX Surgeon since it's kind of a niche question. If it goes against rules, I understand if you delete it, but any help is more than welcome!

Start with the embedded ONNX GraphSurgeon script and the simple inference script, then inspect the generated knn-fixed-2.onnx graph and the Slice inputs reported by ONNX Runtime. Confirm the tensor shapes and determine whether the issue belongs in the model construction or runtime usage; the issue does not define a repository file, test, or acceptance criterion for done.