Perform element-wise multiplication along the columns

Hi, I am trying to multiply two tensors of size (N, d) and (N, d) in the sense that each row is multiplied element-wise and summed over. What is the best practice?

More specifically, suppose that the two tensors are (e.g., N = 2, d = 4)

A = torch.Tensor([[1, 2, 3, 4], [5, 6, 7, 8]])
B = torch.Tensor([[4, 3, 2, 1], [8, 7, 6, 5]])

and I want to get the following result,

torch.Tensor([[20], [164]])

Thanks!

c = torch.mul(A,B)
c = [c[0].sum(), c[1].sum()]
c = [tensor(20.), tensor(164.)]
however c becomes a list

Thanks Laura! I tried also something similar

torch.diagonal(torch.mm(A, B))

but it takes too much space, especially when the two tensors are big.

Do you think there is a walk-around?

C = (A * B).sum(dim=1, keepdims=True)
1 Like

Thank you! It works perfectly.