Thanks for the reply @MariosOreo
I’m not sure if ptblck’s suggestion was to do what I did, but this seemed to work:
I checked the shape of my label tensor, and it was (1,2,512,512)
So in order to match the pos_weight vector’s shape to the label’s shape, I did the following:
pw = torch.FloatTensor([21.956,0.04554])
pw = torch.reshape(pw,(1,2,1,1))
criterion = nn.BCEWithLogitsLoss(pos_weight=pw)

With this in mind, I think the documentation should say something like:
pos_weight  – a weight of positive examples. Must be a vector with length equal to the number of classes along the same dimension as the classes.
Stuff like this might seem obvious to some people, but for newbies like me it really hurts not to have that extra bit of information.
Thanks for the help!