This allows for generic logic which can be used to accomodate arbitrary input/ label tensor partitions.
Currently we treat the SplitNN as a 1d array of models. This allows us to perform horizontal splits in the model.
When this works, we should be able to use the same, standard class for any data/ label distribution