pytorch求矩阵的迹

在PyTorch中，可以使用`torch.trace()`函数来计算矩阵的迹。以下是一个示例：

import torch

# 创建一个2x2的矩阵
matrix = torch.tensor([[1, 2],
                       [3, 4]])

# 计算矩阵的迹
trace = torch.trace(matrix)

print(trace)  # 输出结果: tensor(5)

在这个示例中，我们创建了一个2x2的矩阵`matrix`，然后使用`torch.trace()`函数计算了它的迹，并将结果存储在变量`trace`中。最后，我们打印出了迹的值。