Pytorch中取一个batch的对角线元素

torch.diagonal

这里的j_matrix 是一个（5, 16, 16) 张量

torch.diagonal(j_matrix, dim1=-2, dim2=-1)
Out[6]: 
tensor([[0.3009, 0.7826, 0.2639, 0.4019, 0.5106, 0.3767, 0.4053, 0.2876, 0.3009,
         0.7826, 0.2639, 0.4019, 0.5106, 0.3767, 0.4053, 0.2876],
        [0.8716, 0.1679, 0.3594, 0.5889, 0.4564, 0.9954, 0.3841, 0.5640, 0.8716,
         0.1679, 0.3594, 0.5889, 0.4564, 0.9954, 0.3841, 0.5640],
        [0.5112, 0.2634, 0.4914, 0.8166, 0.3232, 0.4071, 0.4868, 0.7273, 0.5112,
         0.2634, 0.4914, 0.8166, 0.3232, 0.4071, 0.4868, 0.7273],
        [0.4422, 0.3853, 0.6786, 0.3912, 0.5814, 0.3816, 0.7245, 0.2994, 0.4422,
         0.3853, 0.6786, 0.3912, 0.5814, 0.3816, 0.7245, 0.2994],
        [0.2956, 0.4321, 0.9307, 0.5550, 0.5469, 0.6162, 0.5656, 0.2380, 0.2956,
         0.4321, 0.9307, 0.5550, 0.5469, 0.6162, 0.5656, 0.2380]])

标签：0.4019,torch,0.4053,0.5469,0.8166,0.3594,batch,Pytorch,中取
来源： https://blog.csdn.net/weixin_43718238/article/details/113807259