torch.eye

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

Возвращает двумерный тензор с единицами на диагонали и нулями в остальных позициях.

Параметры:

• 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.

Возвращает:

Двумерный тензор с единицами на диагонали и нулями в остальных позициях

Тип возвращаемого значения:

Tensor

Пример:

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