import torch


def argmax_2D(cmaps: torch.Tensor):
    """Retrieve 2D argmax coordinates from dense cost maps.
    Args:
        cmaps: A [N x H x W] tensor.
    Returns:
        argmax_2d: The [N x 2] 2D argmax coordinates.
    """
    channels, _, width = cmaps.shape
    _, argmax_1d = cmaps.view(channels, -1).max(1)
    x_indices = (argmax_1d % width).view(-1, 1)
    y_indices = (argmax_1d // width).view(-1, 1)
    return torch.cat((x_indices, y_indices), dim=1)