Hi Samuel!



 Samue1:

Is there some torch.diag() function that also works for batches?


torch.diagonal() 1 does what you want:
>>> import torch
>>> torch.__version__
'1.9.0'
>>> T = torch.tensor (
...     [[[0.9527, 0.8821],
...     [0.8442, 0.6147]],
...     [[0.0672, 0.5737],
...     [0.1963, 0.4532]],
...     [[0.0992, 0.3838],
...     [0.4169, 0.0925]]]
... )
>>> torch.diagonal (T, dim1 = -2, dim2 = -1)
tensor([[0.9527, 0.6147],
        [0.0672, 0.4532],
        [0.0992, 0.0925]])

Best.
K. Frank