Solved by KFrank in post #2 
              
              
                Hello Zeyuyun! 

The general-purpose tool for taking a product of (contracting) multiple 
tensors along various axes is torch.einsum() (named after “Einstein 
summation”). 
(You can also fiddle with the dimensions to get them to line up as 
needed and use matmul() or bmm().) 
Here is a script that c…