torch.Tensor.size

Tensor.size(dim=None) → torch.Size or int

Возвращает размер тензора self. Если dim не указан, возвращаемое значение имеет тип torch.Size, подкласс tuple. Если dim указан, возвращает целое число, содержащее размер этого измерения.

Параметры:

dim (int, необязательно) – Измерение, размер которого необходимо получить.

Пример:

>>> t = torch.empty(3, 4, 5)
>>> t.size()
torch.Size([3, 4, 5])
>>> t.size(dim=1)
4