Solved by futscdav in post #2 
              
              
                Don’t need a mm operation for that, use broadcasting 
batch_size = 128
a = torch.randn(batch_size, 3, 3)
b = torch.randn(batch_size, 1)

c = a * b.view((-1, 1, 1))