Does this code have implementations for weighted loss?

[Mahi-Mai]

It's my understanding that one of the keys to success for the first uNet was the implementation of weighted loss. This is something I've always struggled to implement myself- my holy grail- and I'm curious if this repo or anyone else have managed to do it.

Here's a screenshot from the original paper showing what I mean:

Thanks so much for this repo!