import torch


def calc_weighted_average(vals, weight):
    """Calculate weighted average along the second dim of values

    Args:
        vals (torch.Tensor): The values to weight; the first dim is samples
        weight (torch.Tensor): The 1d weights to apply to the second dim of vals

    Returns:
        result (torch.Tensor): The result

    """
    weight = weight[None, ...].repeat([vals.shape[0], 1])
    result = torch.mean(weight * vals)
    return result