Ouch, supporting backprop through such a “buffer” is tough, I don’t think pytorch has a solution for that. Maybe writing your “gatherings” inside autograd.Function could work, then in backward function you’d have to scatter gradients to original positions in v (torch.gather does backward with scatter_add_, but your case is more complex, so you’d have to figure how to map indexes).