import torch


def average_losses(all_losses):
    """Average the losses into one dict of losses.
    Args:
        all_losses: List of dictionary of losses.
    Returns:
        combined: A dictionary with same keys as individual dicts, with
            all losses combined.
    """
    if len(all_losses) == 0:
        return {}
    combined = {}
    for key, val in all_losses[0].items():
        if not isinstance(val, torch.Tensor):
            # If it's none or sth.. eg some loss was not active
            combined[key] = val
        else:
            # Average all the values
            stkd = torch.stack([el[key] for el in all_losses])
            # Average the losses that are positive, since I set undefined
            # losses to -1 (where not enough GT is available, etc)
            combined[key] = torch.mean(stkd * (stkd >= 0), dim=0)
    return combined