def get_loss_weights(config: dict):
    """
    Create the correct loss weights dict.
    """

    # Read config file
    mode = config['mode']
    main_out_weight = config['main_out_weight']
    a_out_weight = config['a_out_weight']
    v_out_weight = config['v_out_weight']

    weights = {}
    weights['main_out'] = main_out_weight

    if mode == 's_av':
        weights['a_out'] = a_out_weight
        weights['v_out'] = v_out_weight

    return weights