import torch
import torch.nn as nn
# import torch.nn.functional as F
# import time, pdb, numpy
from utils import accuracy


class LossFunction(nn.Module):
    def __init__(self, nOut, nClasses, **kwargs):
        super(LossFunction, self).__init__()

        self.test_normalize = True

        self.criterion = torch.nn.CrossEntropyLoss()
        self.fc = nn.Linear(nOut, nClasses)

        # print('Initialised Softmax Loss')

    def forward(self, x, label=None):
        x = self.fc(x)
        nloss = self.criterion(x, label)
        prec1, _ = accuracy(x.detach().cpu(),
                            label.detach().cpu(),
                            topk=(1, 5))

        return nloss, prec1
