Sending Input Image Channel to Loss Function

[rahulghosh2]

Hello,

I am trying to implement a new loss function into nnUNet. For this loss function, one of the inputs is the the third channel of the input. For example, my input is 2D-images with 3 channels, resulting in an input tensor shape of is (1, 3, 512, 512). I want the third channel if each input image to be used as an input to the loss function.

I know that I have to define a new trainer class for a new loss function, but I am lost on how to send one of my input channels to the loss function.

Thank you in advance for any advice.