import torch


def softmax(correlation_maps: torch.Tensor):
    """Applies a 2D spatial softmax operation on the dense correlation maps.
    Args:
        * correlation_maps: The batch of the dense correlation maps, of size [N x H x W]
    Returns:
        * The correspondence maps.
    """
    batch, height, width = correlation_maps.shape
    return correlation_maps.view(batch, -1).softmax(dim=1).view(batch, height, width)