pytorch 生成单位张量torch.eye

torch.eye(n, m=None, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor

Returns a 2-D tensor with ones on the diagonal and zeros elsewhere.

>>> torch.eye(3)
tensor([[ 1.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  1.]])