On second thought, in light of the transforms being used, a per pixel pos_weight may have some unintended results. For example, if a given pixel had no positive ground truths in the entire set, it would be divide by zero(quite plausible given this set). That could be dealt with by setting pos_weight[torch.isinf(pos_weight)|torch.isnan(pos_weight)] = 124.26.
Another example, if a given pixel had only one positive ground truth, so pos_weight=58/1 there, but due to the transforms, it ends up seeing 2 or 3 positive ground truths, that might unbalance the training, as well.
To deal with this, you could first loop through that class-specific adding function 100 or so times on all the training set, but while applying the transforms you intend to apply, before dividing for the pos_weight tensor. That would have the benefit of statistically capturing what will likely occur during training. And then just apply the isinf filter to set any “dead” pixels to the mean pos_weight you found earlier.