Skip to content

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")