c = torch.tensor([[1.0000, 2.0000],
                      [0.0000, 0.0000],
                      [0.0000, 0.0000]])

    print(torch.sum(torch.pow((c), 2), axis=0))

tensor([1., 4.])

    c = torch.tensor([[1.0000, 2.0000],
                      [0.0000, 0.0000],
                      [0.0000, 0.0000]])

    print(torch.sum(torch.pow((c), 2), axis=-1))
tensor([5., 0., 0.])

    c = torch.tensor([[1.0000, 2.0000],
                      [0.0000, 0.0000],
                      [0.0000, 0.0000]])

    print(torch.sum(torch.pow((c), 2)))
tensor(5.)