class custom_loss(nn.Module):
    def __init__(self):
        super(custom_loss, self).__init__()

    def forward(self, output, target):
        return F.cross_entropy(output, target)