import torch


def compute_loss(inputs, outputs, criterion, edge_criterion):
    """Compute loss automatically based on what the model output dict contains.
    'doc_logits' -> document-label CE loss
    'para_logits' -> paragraph-label CE loss
    'pare_edge_weights' -> additional edge CE loss
    """
    if 'para_logits' in outputs:
        loss = criterion(outputs['para_logits'], inputs['para_label'])
        loss *= (inputs['para_lengths'] != 0).to(torch.float)
        loss = loss.sum()
    else:
        loss = criterion(outputs['doc_logits'], inputs['label']).sum()

    if 'para_edge_weights' in outputs:
        edge_loss_lambda = 0.1
        edge_loss = edge_criterion(outputs['para_edge_weights'], 1 - inputs['para_label'])
        edge_loss *= (inputs['para_lengths'] != 0).to(torch.float)
        loss += edge_loss_lambda * edge_loss.sum()

    return loss