Function apply_reduction

losses/loss_utils.py:4–10

Apply reduction to collection of losses.

(losses, reduction="none")

import torch

def apply_reduction(losses, reduction="none"):
    """Apply reduction to collection of losses."""
    if reduction == "mean":
        losses = losses.mean()
    elif reduction == "sum":
        losses = losses.sum()
    return losses