Split bug forward dims
Required prerequisites
-
Make sure you've read the documentation. Your issue may be addressed there. -
Search the issue tracker and discussions to verify that this hasn't already been reported. +1 or comment there if it has.
What commit version of aidge do you use
-
aidge_core
: 0.2.2
Problem description
Split
with 3 outputs and the split attribute provided as an input of the operator (3 here)
Assertion failed: splits.size() == nbOutput in /aidge/aidge/aidge_core/src/operator/Split.cpp:109
Split_Op: number of slices [384, 384, 384, 384, 384, 384] must be equal to number of outputs 3
Reproducible example code
#generate ONNX
import torch
import torchvision
# Define a simple model that splits the input
class SplitModel(torch.nn.Module):
def __init__(self, split_size):
super(SplitModel, self).__init__()
self.split_size = split_size
def forward(self, x):
return torch.split(x, self.split_size, dim=-1)
# Create an instance of the model
model = SplitModel(split_size=int(1152/3))
# Define an example input
example = torch.randn(1, 1, 257, 1152)
# Export the model to ONNX
torch.onnx.export(model, # model being run
example, # model input (or a tuple for multiple inputs)
"split_model.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=12, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes
'output' : {0 : 'batch_size'}})
print("Model has been converted to ONNX")