快捷方式

torch.eye

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

返回一个 2-D 张量,对角线上为 1,其余地方为 0。

参数
  • n (int) – 行数

  • m (int, 可选) – 列数,默认为 n

关键字参数
  • out (Tensor, 可选) – 输出张量。

  • dtype (torch.dtype, 可选) – 返回张量的期望数据类型。默认值:如果 None,则使用全局默认值(请参阅 torch.set_default_dtype())。

  • layout (torch.layout, 可选) – 返回张量的期望布局。默认值:torch.strided

  • device (torch.device, 可选) – 返回张量的期望设备。默认值:如果 None,则对默认张量类型使用当前设备(请参阅 torch.set_default_device())。 device 对于 CPU 张量类型将为 CPU,对于 CUDA 张量类型将为当前 CUDA 设备。

  • requires_grad (bool, 可选) – 如果 autograd 应该记录在返回的张量上的操作。默认值:False

返回

一个 2-D 张量,对角线上为 1,其余地方为 0

返回类型

Tensor

示例

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源