You could use a weighted loss e.g. by using an unreduced loss (reduction='none'), multiplying each loss “pixel” by the class weight, reducing it, and finally calling backward().
I’m sure there are other valid approaches. E.g. is focal loss only used in detection models or could it also be used for segmentation use cases?