import torch


def update_weights(batch, batch_dic):
    """
    Readjust weights so they sum to 1.
    Args:
        batch_dic (dict): Dictionary with extra conformer 
            information about the batch
        batch (dict): Batch dictionary
    Returns:
        new_weights (torch.Tensor): renormalized weights
    """

    old_weights = batch["weights"]

    conf_idx = torch.LongTensor(batch_dic["conf_idx"])
    new_weights = old_weights[conf_idx]
    new_weights /= new_weights.sum()
    if torch.isnan(new_weights).any():
        new_weights = torch.ones_like(old_weights[conf_idx])
        new_weights /= new_weights.sum()
    return new_weights