import torch


def squash(x, dim=2):
    """ Computes the squash function of a Tensor along dimension dim

    Args:
        x: input tensor
        dim: dimension along which the squash must be performed

    Returns:
        A tensor squashed along dimension dim of the same shape of x """
    norm = torch.norm(x, dim=dim, keepdim=True)
    return x * norm / (1 + norm ** 2)