Need help in weighted-loss calculation

I am working on a tumor segmentation problem and need help in calculating the weighted loss in BCEWithLogitsLoss function.
So, I am having image (240,240,155) and mask (240,240,155) with labels - 0,1,2,4

Now in training i have image (1,4,144,240,240) and mask (1,3,144,240,240) shape with labels 0,1
How to calculate the weighted loss here.