Custom, weighted loss function for segmentation tasks - Error: object has no attribute '_forward_pre_hooks'

I’m working on a segmentation project, where I tried to segment 14 classes of objects inside an image. Each image consists mostly of background pixels, while the rest of the pixels belong to some of the 14 classes of the objects I’m trying to segment.

What I’m trying to do is pass on weights to the loss function, as in here. Further discussion is here. Passing those weights to the loss function means the gradients of background pixels are weighted down compared to the pixels belonging to the 14 objects, to avoid having the output of the network converge to zero.

I’m getting a AttributeError: 'CrossEntropyFlatWeighted' object has no attribute '_forward_pre_hooks' error and it happens in` file in the torch code.

This screenshot of my Notebook shows what I’m trying to do and the error I’m getting! If anyone did something similar successfully I’d appreciate your help!


This way is working for me, just use the default CrossEntropyFlat model. They allow arguments for nn.CrossEntropyLoss to take.

learn.loss_fn = CrossEntropyFlat(weight=class_weights)