I have a very large n x n tensor and I want to fill its diagonal values to zero, granting backwardness. How can it be done? Currently the solution I have in mind is this

However if n is large this can potentially require a lot of memory. Is there a simpler differentiable solution, perhaps similar to NumPy np.fill_diagonal function?

How to do it for a batch of tensors? For tensor with shape N * M * M. Have to set diagonal equal to zero for every M*M matrix. fill_diagonal doesn’t work if N isn’t equal to M.