Solved by tom in post #3 
              
              
                For two classes, the usual setup is to take the binary cross entropy loss. PyTorch has nn.BCELoss and nn.BCELossWithLogits for this, the latter including the sigmoid layer that you would need to add to your net for the former. The trick here is that the “probability” of the 0 class is just 1 minus t…