melike:

I wrote it before learning that in-place operations should be avoided in PyTorch.


You don’t have to avoid them. It is just that autograd does not support every combination of them and it will raise an error if you hit such case.
So if your code runs without error, it means that autograd can handle this case just fine.
The only concern I would have with such implementation is the slowdown due to the nested loops. But that’s unrelated to gradient correctness.