import torch


def squash(tensor, dim=-1):
    """ Squash function as defined in [1].

    Args:
        tensor (Tensor): Input tensor.
        dim (int, optional): Dimension on which to apply the squash function. Vector dimension. Defaults to the last.
    """
    squared_norm = (tensor ** 2).sum(dim=dim, keepdim=True)
    scale = squared_norm / (1. + squared_norm)
    return scale * tensor / torch.sqrt(squared_norm + 1e-7)