squeeze的源代码

以下是squeeze函数的源代码，该函数用于将张量的维度压缩到指定大小：

def squeeze(input, dim=None, out=None):
    """
    Returns a tensor with all the dimensions of :attr:`input` of size `1` removed.

    For example, if `input` is of shape:

    .. code-block:: python

        input.size() = (A x 1 x B x C x 1 x D)

    Then the `output` tensor will be of shape:

    .. code-block:: python

        output.size() = (A x B x C x D)

    When :attr:`dim` is given, a squeeze operation is done only in the given
    dimension. If `input` is of shape: 

    .. code-block:: python

        input.size() = (A x 1 x B)

    `squeeze(input, 0)` leaves the tensor unchanged, but `squeeze(input, 1)` will
    squeeze the tensor to the shape `(A x B)`.

    .. note:: The returned tensor shares the storage with the input tensor,
              so changing the contents of one will change the contents of the other.

    Args:
        input (Tensor): the input tensor
        dim (int, optional): if given, the input will be squeezed only in
            this dimension
        out (Tensor, optional): the output tensor

    Example::

        >>> x = torch.zeros([1, 3, 1, 2, 1])
        >>> x.size()
        torch.Size([1, 3, 1, 2, 1])
        >>> y = torch.squeeze(x)
        >>> y.size()
        torch.Size([3, 2])
        >>> y = torch.squeeze(x, 0)
        >>> y.size()
        torch.Size([3, 1, 2, 1])
    """
    if out is not None:
        return torch._C._VariableFunctions.squeeze(input, dim, out)
    return torch._C._VariableFunctions.squeeze(input, dim)