📂머신러닝

파이토치에서 텐서 차원 유지하면서 계산하는 방법

코드

함수를 쓸 때 옵션으로 keepdim = True를 넣어주면 된다.

>>> A
tensor([[1., 1., 1., 1.],
        [2., 2., 2., 2.],
        [3., 3., 3., 3.],
        [4., 4., 4., 4.]])

>>> A.sum(dim=1)
tensor([ 4.,  8., 12., 16.])

>>> A.sum(dim=1, keepdim=True)
tensor([[ 4.],
        [ 8.],
        [12.],
        [16.]])

>>> torch.linalg.norm(A, dim=1)
tensor([2., 4., 6., 8.])

>>> torch.linalg.norm(A, dim=1, keepdim=True)
tensor([[2.],
        [4.],
        [6.],
        [8.]])

환경

• OS: Windows11
• Version: Python 3.10.11, torch 2.7.0+cu126