nn.BCEWithLogitsLoss expects the model output to contain logits, which are not bounded to a specific range.
Based on your code it seems that this loss calculation is correct (I haven’t checked the kld loss).
However, since your target image contains zeros and ones, I would assume that you would need to apply torch.sigmoid on the model output before storing the predictions.
Otherwise negative logits would correspond to a low probability (close to 0) while positive logits would correspond to a high probability (close to 1).
You might also or alternatively apply a threshold to create only zeros and ones as the predicted image.