losses/volume_label_loss.py

import torch.nn as nn


def get_volume_label_loss(opt):
    return VolumeLabelLoss()


class VolumeLabelLoss(nn.Module):
    def __init__(self):
        super().__init__()

        self.BCE_loss = nn.BCELoss(reduction="mean")

    def forward(self, pred, volume, label):
        loss = self.BCE_loss(pred, label)
        return {"loss": loss}