m = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.]])   # shape (2, 3)

print(m.sum())          # tensor(21.)          — no dim: collapse everything to a scalar
print(m.sum(dim=0))     # tensor([5., 7., 9.]) — collapse the rows, shape (3,)
print(m.sum(dim=1))     # tensor([ 6., 15.])   — collapse the columns, shape (2,)