PyTorch torch.split 函数
Pytorch torch 参考手册torch.split 是 PyTorch 中用于将张量沿指定维度分割的函数。
函数定义
torch.split(tensor, split_size, dim)
使用示例
实例
import torch
x = torch.arange(12).reshape(3, 4)
# 沿 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)
# 沿 dim=0 分割
result = torch.split(x, 1, dim=0)
for i, t in enumerate(result):
print(f"Part {i}:")
print(t)