如何在 4d 火炬张量中获取 3d 张量模块？

Question

I have a torch tensor with 4 dimensions. It's shape is [50, 1, 1, 200]. I have to get a list with modules of 200 hundred 3d tensors. Which is the easiest way to do so?

Answer 1

Did you try torch.unbind ( https://pytorch.org/docs/stable/generated/torch.unbind.html )?

a = torch.rand(50, 1, 1, 200)
b = torch.unbind(dim=3)

len(b)      # 200
b[0].shape  # torch.Size([50, 1, 1])