def irm_loss_rescale(total_loss, irm_penalty_weight):
    """
    Rescale the total loss by the IRM penalty weight as done by
    https://github.com/facebookresearch/InvariantRiskMinimization
    """
    if irm_penalty_weight > 1.0:
        return total_loss / irm_penalty_weight
    return total_loss