torch.split
torch.split(tensor, split_size_or_sections, dim=0)
Splits the tensor into chunks. Each chunk is a view of the original
   tensor.
If "split_size_or_sections" is an integer type, then "tensor" will
   be split into equally sized chunks (if possible). Last chunk will
   be smaller if the tensor size along the given dimension "dim" is
   not divisible by "split_size".
If "split_size_or_sections" is a list, then "tensor" will be split
   into "len(split_size_or_sections)" chunks with sizes in "dim"
   according to "split_size_or_sections".
Parameters:
      * tensor (Tensor) -- tensor to split.
  * **split_size_or_sections** (*int**) or **(**list**(**int**)*)
    -- size of a single chunk or list of sizes for each chunk

  * **dim** (*int*) -- dimension along which to split the tensor.

Return type:
      List[Tensor]
Example:
  >>> a = torch.arange(10).reshape(5, 2)
  >>> a
  tensor([[0, 1],
          [2, 3],
