torch.numel

torch.numel

input tensor에 대해 element의 총 개수를 return

>>> a = torch.randn(1, 2, 3, 4, 5)
>>> torch.numel(a)
120
>>> a = torch.zeros(4,4)
>>> torch.numel(a)
16

'Programming > Pytorch' 카테고리의 다른 글

Custom mini-batch Loader  (0) 2023.03.21