I have a tensor which has some diagonal matrices inside it.

```
tensor([[[-0.2920, 0.0000],
[ 0.0000, -0.8663]],
[[-0.2920, 0.0000],
[ 0.0000, -0.7568]],
[[-0.3373, 0.0000],
[ 0.0000, -0.5953]],
[[-0.0068, 0.0000],
[ 0.0000, -0.8065]],
[[-0.1260, 0.0000],
[ 0.0000, -0.7147]]], grad_fn=<MulBackward0>)
```

I want to do an operation on this to change this tensor to the following such that I do not lose the gradients because I need to call backward later.

```
tensor([[-0.2920, -0.8663],
[-0.2920, -0.7568],
[-0.3373, -0.5953],
[-0.0068, -0.8065],
[-0.1260, -0.7147]])
```

I saw `diagonal_embed`

but it didnt help much. How do I do this?