Usage of squeeze function in pytorch

The Chinese meaning of squeeze is " squeeze ". As the name suggests, the function of this function is to compress the dimensions.

squeeze(input, dim=None) -> Tensor

Input a high-dimensional tensor. Squeeze will only work if there is a dimension of size 1 in each dimension. Here is an example.

x = torch.arange(6).reshape(2,1,3)

# tensor([[[0, 1, 2]],
#         [[3, 4, 5]]])  shape=(2,1,3)

x = x.squeeze()

# tensor([[0, 1, 2],
#        [3, 4, 5]])    shape=(2,3)

If dim is specified, when dim=1, the effect is the same as above. If dim is other, then the dimension of x remains unchanged.

Of course, if there is no dimension of size 1 in each dimension, the squeeze function has no effect on the x tensor.