Dear altruists,

I have two tensors with dim (2, 4, 4)) and (4,4) respectively.

A = tensor([[[0, 0, 1, 0],

[0, 0, 1, 0],

[1, 1, 0, 1],

[0, 0, 1, 0]],

[[0, 0, 1, 0],

[0, 0, 1, 0],

[1, 1, 0, 1],

[0, 0, 1, 0]]])

B = tensor([[0, 1, 1, 0],

[1, 0, 1, 0],

[1, 1, 0, 1],

[0, 0, 1, 0]])

I want to do multiplication along second and third axes. My desired output is a vector with dim 2.

I am trying to take advantage of the speed of torch.cuda() for matrix multiplication. Is there an elegant way to solve this problem?

Thanks in advance!