Note that F.cross_entropy does not need a softmax, but expects the raw logits. You should therefore not use any activation function after your last linear layer. (Unsure, if you are going to swap the sigmoid for a softmax)