Hdk:

Now, I have to include an additional loss in the backward computation of the layer itself


Could you explain why you have to include this in the backward of the layer itself?
Why can’t it be an additional term of your loss that you backprop as usual (ce_loss + additional_loss).backward() ?