Solved by vainaijr in post #2 
              
              
                do you mean something like this, 
import torch
x = torch.randn(2, 14, 14, 128)
y = torch.randn(2, 128)
torch.einsum('abcd, ad -> abc', x, y).unsqueeze(-1).shape

torch.Size([2, 14, 14, 1])