Setting diagonal terms of a tensor to zero

I have a symmetric tensor of size [B, N, N, C]. I would like to eliminate the diagonal elements. I was wondering what is the most efficient way to do this?

My solution:


        idx1 = torch.triu_indices(N, N, 0)
        idx2 = torch.triu_indices(N, N, 1)

Then turn them into lists, subtract to get diagonal indices. Then use them to subtract diagonal term from the original tensor. But I’m wondering

  1. Would turning indices into list break backprop?
  2. Also, I’m not sure if this is the most efficient way to do this.

Interesting, does it have to be inplace?
How about creating an inverted identity diagonal matrix, and then multiply element wise?
your_tensor *= (1 - torch.eye(N, N))

Also checkout fill_diagonal_ function it might fit perfectly


1 Like