Tensor.tensor_split()

Tensor.tensor_split(indices_or_sections, dim=0) → List of Tensors

请参阅 torch.tensor_split()