如果您处理一个较小的示例，就可以理解这个问题：

A = torch.randint(2, (4,))
B = torch.eye(2)

>>> A
# tensor([1, 0, 1, 1])

>>> B[A].shape
# (4, 2)

>>> B[A]
# tensor([[0., 1.],
#         [1., 0.],
#         [0., 1.],
#         [0., 1.]])

[1, 0]和[0, 1]是2x2单位矩阵B的第一行和第二行。因此，使用形状为（4，）的一维数组A作为索引是选择B的4个“行”/沿着轴0选择B的4个元素。B[A]基本上是[B[1], B[1], B[0], B[1]]。
因此，当A是形状为(1, 4, 4)的三维数组时，B[A]意味着**选择B的（1，4，4）行。**由于B中的每一行都有2个元素（2列），因此输出为（1，4，4，2）。
B是一个2x2单位矩阵，有2行。从这两行中选出16行，得到一个（16，2）矩阵，然后将其整形得到（1，4，4，2）Tensor。实际上，你可以很容易地检查：

A = torch.randint(2, (4, 4))
A_flat = A.reshape(-1)
B = torch.eye(2)

>>> torch.allclose(B[A], B[A_flat].reshape(1, 4, 4, -1)])
# True

这也不是PyTorch特有的现象，你可以在NumPy中观察到同样的索引规则，它与torch保持着紧密的兼容性。