yash1994:

nn.NLLLoss and nn.BCELoss loss function requires softmax to be applied beforehand. But any of these loss functions don’t require the following operations you’ve performed (because these operations are implemented internally).


Hi @yash1994
Thank you so much for your response.
I am using nn.BCEWithLogitsLoss as my loss function. As you had suggested, I removed the softmax operation from my model.forward function and the two lines related to predictions that you had highlightened.
I receive the following error
Value Error: Target size (torch.Size[64)) must be the same as the input size (torch.Size([64,2]))
My updated model training step is as below and I’m guessing the size of predictions is [64,2] while the size of batch.label is [64]
# #Reset gradients after every batch
        optimizer.zero_grad()
        
        predictions = model(batch.X)
       
        loss = criterion(predictions, batch.label)
        loss.backward()  # #Error Here

        optimizer.step()

Any suggestions on resolving this?