Yes, I’m using binary classification. But with the code I provided above, I get a probability distribution over the 2 classes I have, and my final layer is already a nn.Linear(1024, 2), but I train the network with a crossentropy criterion… My doubt is if make sense to add a softmax on top of the output which is a result of a crossentropy loss.