from torch import nn


class ColorLoss(nn.Module):
    def __init__(self, coef=1):
        super().__init__()
        self.coef = coef
        self.loss = nn.MSELoss(reduction="mean")

    def forward(self, inputs, targets):
        loss = self.loss(inputs["rgb_coarse"], targets)
        if "rgb_fine" in inputs:
            loss += self.loss(inputs["rgb_fine"], targets)

        return self.coef * loss


loss_dict = {"color": ColorLoss}
