torch.squeeze

torch.squeeze(input, dim=None, *, out=None) -> Tensor

공식설명은 아래와 같다.

Returns a tensor with all the dimensions of input of size 1 removed.

For example, if input is of shape: (A \times 1 \times B \times C \times 1 \times D)(A×1×B×C×1×D) then the out tensor will be of shape: (A \times B \times C \times D)(A×B×C×D).

When dim is given, a squeeze operation is done only in the given dimension. If input is of shape: (A \times 1 \times B)(A×1×B), squeeze(input, 0) leaves the tensor unchanged, but squeeze(input, 1) will squeeze the tensor to the shape (A \times B)(A×B).

input의 dimension에서 size 1이 제거된 tensor를 반환한다. 예를 들면, input의 shape이 (A x 1 x B x C x 1 x D) 일 때 squeeze를 사용하면 output tensor의 shape이 (A x B x C x D)로 변경된다.

dim argument가 주어졌을 때 squeeze연산은 주어진 dimension에 대해서만 수행된다. 예를 들면 input shape이 (A x 1 x B)일 때, squeeze(input, 0)은 tensor를 변화시키지 못한다. squeeze(input, 1)을 이용하면 output tensor가 (A x B) shape을 갖게 된다.

Reference

1. https://pytorch.org/docs/stable/generated/torch.squeeze.html