TermoPile:

however, this operation is only implemented for 2D tensors


That’s not the case, as you can use tensors with having multiple dimensions with index_add_.
I don’t quite understand your use case, as it seems you would like to add the tensor to itself instead of another passed value tensor?
Could you post an example using randomly initialized tensors?