import torch


def compute_argmax(ten):
    """Compute argmax for 2D grid for tensors of shape (batch_size, size_y, size_x)

    Args:
        ten (torch.[cuda].FloatTensor): (batch_size, size_y, size_x)
    Returns:
        indices (torch.[cuda].LongTensor): (batch_size, 2) index order: (y, x)
    """

    batch_size = ten.shape[0]
    size_y = ten.shape[1]
    size_x = ten.shape[2]

    # shape: flattened_indices (batch_size)
    flattened_indices = torch.argmax(ten.view(batch_size, -1), dim=1)

    # shape: index_y (batch_size)
    # shape: index_x (batch_size)
    index_y = (flattened_indices // size_x)
    index_x = (flattened_indices % size_x)

    # shape: index_y (batch_size, 2)
    indices = torch.cat([index_y.unsqueeze(1), index_x.unsqueeze(1)], dim=1)

    return indices