Hi All,

I’m trying to figure out a way to set the diagonal of a 3-dimensional Tensor (along 2 given dims) equal to 0. An example of this would be, let’s say I have a Tensor of shape `[N,N,N]`

and I wanted to set the diagonal along dim=1,2 equal to 0? How exactly could that be done?

I tried using `fill_diagonal_`

but that only does the k-th diagonal element for each sub-array, i.e,

```
data = torch.ones(4,4,4)
data.fill_diagonal_(0)
tensor([[[0., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]],
[[1., 1., 1., 1.],
[1., 0., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]],
[[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 0., 1.],
[1., 1., 1., 1.]],
[[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 0.]]])
```

whereas I would want the entire diagonal for each sub-matrix to be equal to 0 here. So, the desired outcome would be,

```
tensor([[[0., 1., 1., 1.],
[1., 0., 1., 1.],
[1., 1., 0., 1.],
[1., 1., 1., 0.]],
[[0., 1., 1., 1.],
[1., 0., 1., 1.],
[1., 1., 0., 1.],
[1., 1., 1., 0.]],
[[0., 1., 1., 1.],
[1., 0., 1., 1.],
[1., 1., 0., 1.],
[1., 1., 1., 0.]],
[[0., 1., 1., 1.],
[1., 0., 1., 1.],
[1., 1., 0., 1.],
[1., 1., 1., 0.]]])
```

Secondly, the reason I state for a given pair of dimension is, I need to repeat this `zeroing’ along 2 different pairs of dimensions (e.g. dim=(1,2) then dim=(0,1) ) to get the required masking I need.

In short, is there a way to mask a given diagonal over 2 arbitrary dimensions for a 3D-tensor?

Thank you in advance!