Divide a batch into smaller groups
Mini-batch preparation
PyTorch · Function
torch.split
Splits a tensor into chunks along a selected dimension. Use an integer for repeated chunk sizes or a list for explicit sizes; returned chunks are views of the input.
How it works
torch.split divides a tensor along a chosen dimension and returns the resulting pieces as a tuple. An integer requests equal-sized chunks where possible, with a smaller final chunk when necessary; a list specifies the size of every chunk. The outputs are views into the original tensor.
Syntax
torch.split(tensor, split_size_or_sections, dim=0)Parameters
| Name | Type | Description |
|---|---|---|
tensor | Tensor | Tensor to split. |
split_size_or_sections | int or list[int] | Size of one chunk, or a list specifying the size of each chunk. |
dim | int | Dimension along which to split the tensor; defaults to 0. |
Returns
tuple[Tensor, …]
Key points
- The function partitions one tensor into multiple chunks.
- Every returned chunk aliases the original tensor instead of being freshly allocated.
- An integer argument requests equal-sized chunks whenever the dimension permits it.
- A list determines the number of outputs from its length.
- List entries specify each output's extent along the selected dimension.
Common mistakes
- When the selected dimension is not evenly divisible by an integer split size, the final chunk is smaller.
Example: Split a tensor by size or explicit sections
import torch
a = torch.arange(10).reshape(5, 2)
chunks = torch.split(a, 2)
sections = torch.split(a, [1, 4])
assert [tuple(x.shape) for x in chunks] == [(2, 2), (2, 2), (1, 2)]
assert [tuple(x.shape) for x in sections] == [(1, 2), (4, 2)]
The integer form creates chunks of up to two rows, while the list form requests chunks containing one and four rows.
When to use it
Partition time or feature dimensions into windows
Sequence processing
Create explicitly sized sections with a list
Uneven grouping
Interview questions
QWhat does torch.split accept, and what does it return?
It returns a tuple of tensors produced by splitting the input along dim, which defaults to dimension 0. The second argument can be an integer chunk size or a list containing the size of each chunk.
Follow-up: How does an integer split behave when the dimension length is not divisible by the requested size? The final chunk contains the remaining elements and is smaller.
QHow do integer and list split specifications differ?
With an integer split size, PyTorch makes equally sized chunks whenever possible. With a list, it creates one chunk per list entry and uses those entries as the chunk sizes along the selected dimension.
Follow-up: How many outputs does a list such as [1, 4] produce? Exactly two, with sizes one and four along the split dimension.
QAre the returned chunks copies or views?
The outputs are views of the input rather than independent copies. Therefore, they share the input's underlying storage and should be treated as aliases of the original tensor.
Follow-up: Which operation should be considered when independent output storage is required? The documentation lists torch.split_copy as the copy-producing counterpart.
Related
torch.Tensor.splittorch.tensor_splittorch.split_copytorch.split_with_sizes_copy