nn.BCEWithLogitsLoss expects logits as the model output, not probabilities, so you would have to remove the last sigmoid activation in your model. Besides that I don’t see any obvious issues in your code.