如何将Pytorch张量拆分成不同的维度？

import torch

# 创建一个形状为(3, 6)的张量
tensor = torch.arange(18).reshape(3, 6)
print("原始张量：")
print(tensor)

# 拆分张量成3个形状为(2, 6)的子张量
split_tensors = torch.split(tensor, 2)
print("拆分后的张量：")
for split_tensor in split_tensors:
    print(split_tensor)

原始张量：
tensor([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11],
        [12, 13, 14, 15, 16, 17]])
拆分后的张量：
tensor([[0, 1, 2, 3, 4, 5],
        [6, 7, 8, 9, 10, 11]])
tensor([[12, 13, 14, 15, 16, 17]])