PyTorch torch.split function
PyTorch torch Reference Manual
torch.splitIt is a function in PyTorch used to split a tensor along a specified dimension.
Function Definition
torch.split(tensor, split_size, dim)
Usage Example
Example
import torch
x = torch.arange(12).reshape(3, 4)
# Split along dim=0
result = torch.split(x, 1, dim=0)
for i, t in enumerate(result):
print(f"Part {i}:")
print(t)
x = torch.arange(12).reshape(3, 4)
# Split along dim=0
result = torch.split(x, 1, dim=0)
for i, t in enumerate(result):
print(f"Part {i}:")
print(t)
Other Extensions