torch.eye

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

返回一個二維張量，其對角線上是 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()）。對於 CPU 張量型別，device 將是 CPU；對於 CUDA 張量型別，則為當前 CUDA 裝置。

• requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量上的操作。預設值：False。

返回值

一個二維張量，其對角線上是 1，其他地方是 0

返回型別

Tensor

示例

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