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