Did you manage to speed this up? I’m facing the exact same problem. 
Edit 1:
Speeded it up by vectorizing the operations using torch.baddbmm instead of the double for loops.
Used this fact a lot:
For a batch-first scenario,
x = torch.randn(10, 4, 8) # encoder inputs (batch * word * ?)
y = torch.randn(10, 8, 1) # hidden state

# this for loop version
aa = torch.zeros(10, 4)
for batch in range(x.size()[0]):
    for word in range(x.size()[1]):
        aa[batch,word] = x[batch,word].dot(y[batch,:].squeeze(1))

# is equivalent to this vector version
bb = torch.baddbmm(torch.zeros(4,1), x, y).squeeze(2)

# verify if they are equal
aa == bb
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1
    # 1     1     1     1