Weights in BCEWithLogitsLoss