Are there any plans to add an easy way to include "loss mask" to Dataset class for segmentation task? As what I understood, it only takes img and seg?

For example I have such CT volumes with the following "windowing" that could be masked away

And in general the following behavior could be a nice feature

https://github.com/deepmind/deepmind-research/blob/cf0a9ca46058a61e8b66bc8708cd35ac77209bc6/hierarchical_probabilistic_unet/geco_utils.py#L100