torch.Tensor.chunk

Tensor.chunk(chunks, dim=0) → List of Tensors

See torch.chunk()