code: def split(data: mx.nd.NDArray,
            num_outputs: int,
            axis: int = 1,
            squeeze_axis: bool = False) -> List[mx.nd.NDArray]:
      
      ndarray_or_list = data.split(num_outputs=num_outputs, axis=axis, squeeze_axis=squeeze_axis)
      if num_outputs == 1:
          return [ndarray_or_list]
      return ndarray_or_list
docs: Version of mxnet.ndarray.split that always returns a list.  The original
      implementation only returns a list if num_outputs > 1:
      https://mxnet.incubator.apache.org/api/python/ndarray/ndarray.html#mxnet.ndarray.split
  
      Splits an array along a particular axis into multiple sub-arrays.
  
      :param data: The input.
      :param num_outputs: Number of splits. Note that this should evenly divide
                          the length of the axis.
      :param axis: Axis along which to split.
      :param squeeze_axis: If true, Removes the axis with length 1 from the shapes
                           of the output arrays.
      :return: List of NDArrays resulting from the split.