PyTorch torch.eye Function


Pytorch torch 参考手册Pytorch torch Reference Manual

torch.eyeis a function in PyTorch used to create an identity matrix (a matrix with 1s on the diagonal and 0s elsewhere).

This is commonly used in deep learning for initialization, creating masks, and other scenarios.

Function Definition

torch.eye(n, m, dtype, device, requires_grad)

Parameters:

  • n(int): number of rows.
  • m(int, optional): number of columns. If not specified, an n×n square matrix is created.
  • dtype(torch.dtype, optional): data type.
  • device(torch.device, optional): device.
  • requires_grad(bool, optional): whether gradient computation is needed.

Return value:

  • torch.Tensor: returns the identity matrix.

Usage Examples

Example 1: Create a Square Matrix

Example

import torch

# Create a 3x3 identity matrix
I = torch.eye(3)

print(I)

The output is:

tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])

Example 2: Create a Non-Square Matrix

Example

import torch

# Create a 3x4 identity matrix
I = torch.eye(3, 4)

print(I)

The output is:

tensor([[1., 0., 0., 0.],
        [0., 1., 0., 0.],
        [0., 0., 1., 0.]])

Pytorch torch 参考手册Pytorch torch Reference Manual

Other extensions