torch.tensor_split
torch.tensor_split(input, indices_or_sections, dim=0) -> List of Tensors
Splits a tensor into multiple sub-tensors, all of which are views
   of "input", along dimension "dim" according to the indices or
   number of sections specified by "indices_or_sections". This
   function is based on NumPy's "numpy.array_split()".
Parameters:
      * input (Tensor) -- the tensor to split
  * **indices_or_sections** (*Tensor**, **int** or **list** or
    **tuple of ints*) --

    If "indices_or_sections" is an integer "n" or a zero
    dimensional long tensor with value "n", "input" is split into
    "n" sections along dimension "dim". If "input" is divisible by
    "n" along dimension "dim", each section will be of equal size,
    "input.size(dim) / n". If "input" is not divisible by "n", the
    sizes of the first "int(input.size(dim) % n)" sections will
    have size "int(input.size(dim) / n) + 1", and the rest will
