torch.squeeze用法

torch.squeeze：将tensor中大小为1的维度删除
torch.squeeze(tensor,dim）
代码示例：

import torch
a=torch.ones(2,2,2,1,1)
b=torch.squeeze(a)
c=torch.squeeze(a,0)
d=torch.squeeze(a,3)
print("a的shape：",a.shape)
print("b的shape：",b.shape)
print("c的shape：",c.shape)
print("d的shape：",d.shape)

输出：

a的shape： torch.Size([2, 2, 2, 1, 1])
b的shape： torch.Size([2, 2, 2])
c的shape： torch.Size([2, 2, 2, 1, 1])
d的shape： torch.Size([2, 2, 2, 1])

注：对特定维度进行squeeze，如果该维度大小不为1，则保持原来的shape不变，如代码中的c的shape。