import torch


def get_bin_idx_loss(gt_bin_idx, pred_bin_idx, loss_type):
    """
    gt_bin_idx: [B]
    pred_bin_idx: [B]

    Loss: L2 (MSE), or L1
    """
    if loss_type == "L2":
        loss = (gt_bin_idx - pred_bin_idx) ** 2  # [B]
    elif loss_type == "L1":
        loss = torch.abs(gt_bin_idx - pred_bin_idx)  # [B]
    else:
        assert False, "Unknown loss type"

    return loss