Hi,
The backward that you implement is not the same as the one that gets called. The one you write will only work with gradients.
You can solve this by simply passing both the from_rank and to_rank to the apply (and so take both as input for your forward).
And in your forward save anything that you need for the backward on the ctx: ctx.to_rank = to_rank.