The model performance is greatly affected by long-tailed dataset

Sorry to bother you.
I trained Prediff on my own dataset and found the result quite bad. I guess the reason behind should be data imbalance which is commonly observed in precipitation nowcasting.
I am currently considering to do resampling but I am worrying that it might hurt the generalizability.
I noticed that in your previous paper regarding TrajGRU, pixelwise loss weighting is applied to the radar sequence. How could I implement a similar approach in Prediff.
I would be appreciated if you could offer some suggestions.