为何在尺寸为1的维度上使用:和0会得到不同输出？

先看你的张量基础信息：

x = torch.arange(1,10).reshape(3,1,3)

这个张量的形状是(3, 1, 3)，属于三维张量，其中第二个维度的长度仅为1。

两者输出不同的核心原因是是否保留被索引维度的结构：

• x[:,:,2]：第二个维度使用:表示“选取该维度的全部元素”，即便该维度只有1个元素，PyTorch也会完整保留这个维度的结构。最终结果的形状为(3, 1)，对应输出的[[3], [6], [9]]——这是一个二维张量，每个元素都是单元素的子张量。
• x[:,0,2]：第二个维度使用0直接选取该维度的唯一元素，此时PyTorch会自动挤压掉这个长度为1的维度，将原三维张量降为一维张量，形状变为(3,)，对应输出的[3, 6, 9]。

如果想要用索引0但保留维度结构，可以改用切片语法或列表索引，比如x[:,0:1,2]或者x[:,[0],2]，这样得到的结果形状就和x[:,:,2]完全一致了。

内容的提问来源于stack exchange，提问作者Kushagra Mishra