torch.eye

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

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

Параметры

• n (int) – количество строк
• m (int, необязательно) – количество столбцов, по умолчанию n

Ключевые аргументы

• out (Тензор, необязательно) – выходной тензор.
• dtype (torch.dtype, optional) – желаемый тип данных возвращаемого тензора. По умолчанию: если None, используется глобальный по умолчанию (см. torch.set_default_tensor_type()).
• layout (torch.layout, optional) – желаемая структура возвращаемого тензора. По умолчанию: torch.strided.
• device (torch.device, optional) – желаемое устройство возвращаемого тензора. По умолчанию: если None, используется текущее устройство для типа тензора по умолчанию (см. torch.set_default_tensor_type()). device будет ЦПУ для тензоров ЦПУ и текущее CUDA-устройство для тензоров CUDA.
• requires_grad (bool, необязательно) – Если автограду нужно записывать операции над возвращаемым тензором. По умолчанию: False.

Возвращает

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

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

Тензор

Пример:

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