Stone_Yu:

However, I’m still wondering what’s the correct way to use top-k values, as top-k will break the chain rule (not sure just in my opinion)? Could you provide some examples?


topk will propagate gradients from the grad of output values to the grad of inputs all right (in fact, the backward is a scatter to the indices similar to the ones you have), so the topk part you show should work now.
Best regards
Thomas