One thing I see at first is your optimizer.zero_grad() should be happening every step not every epoch. That could be throwing off your gradients. So just add another indent so it is in the dataloader loop not the epoch loop. Also this isn’t a big problem but it is better to use BCEWithLogitsLoss and no sigmoid for training. You can read more about why in the docs here.