torch.split用法

关于torch.split用法简单总结：

torch.split，用来划分tensor，可以从数量上划分，还有维度上划分。
torch.split(tensor,split_szie,dim），split_size有整数，也有列表，dim默认为0，自己也可以修改。

tensor：输入张量

split_size：张量的大小，这里面包含整数或者列表

dim：从哪个维度切割
代码示例1：

x = torch.arange(10).reshape(5, 2)
print(x)
y = torch.split(x, [1, 4])
print(y)

输入是一个5*2维度的张量

对于y的输出是从dim=0进行切割，也就是从5切割，切割的分配比例是第一个tensor维度大小变成1*2；第二个tensor大小变成4*2.

代码示例2：从dim=0，进行分割张量。第一个tensor的shape是：2*2，第二个shape是3*2，

x = torch.arange(10).reshape(5, 2)
print(x)
y = torch.split(x, [2, 3])
print(y)