Hi ixaxaar,
I also have exactly the same problem. So I’m trying to make it faster with baddbmm, though, I have a trouble.
In my case, encoder_outputs has (length x batch_size x hidden_vector_size) and hidden has (1 x batch_size x hidden_vector_size) in forward function. I mean they are not batch first.
class Attn(nn.Module):
    def __init__(self, method, hidden_size):
        super(Attn, self).__init__()

        self.method = method
        self.hidden_size = hidden_size

        if self.method == 'general':
            self.attn = nn.Linear(self.hidden_size, hidden_size)

        elif self.method == 'concat':
            self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
            self.v = nn.Parameter(torch.FloatTensor(1, hidden_size))

    def forward(self, hidden, encoder_outputs):
        max_len = encoder_outputs.size(0)
        this_batch_size = encoder_outputs.size(1)

        # Create variable to store attention energies
        attn_energies = Variable(torch.zeros(this_batch_size, max_len))  # B x S

        if USE_CUDA:
            attn_energies = attn_energies.cuda()

        # For each batch of encoder outputs
        for b in range(this_batch_size):
            # Calculate energy for each encoder output
            for i in range(max_len):
                attn_energies[b, i] = self.score(hidden[:, b], encoder_outputs[i, b].unsqueeze(0))
        #try to replace for loop with baddbmm
        bb = torch.baddbmm(Variable(torch.zeros(max_len, 1).cuda()), encoder_outputs.transpose(0,1), hidden.transpose(0,1)).squeeze(2)
        # Normalize energies to weights in range 0 to 1, resize to 1 x B x S
        return F.softmax(attn_energies).unsqueeze(1)

    def score(self, hidden, encoder_output):

        if self.method == 'dot':
            energy = hidden.dot(encoder_output)
            return energy

        elif self.method == 'general':
            energy = self.attn(encoder_output)
            energy = hidden.dot(energy)
            return energy

        elif self.method == 'concat':
            energy = self.attn(torch.cat((hidden, encoder_output), 1))
            energy = self.v.dot(energy)
            return energy

Instead of for loop, I’m trying to use baddbmm like this.
        for b in range(this_batch_size):
            # Calculate energy for each encoder output
            for i in range(max_len):
                attn_energies[b, i] = self.score(hidden[:, b], encoder_outputs[i, b].unsqueeze(0))


bb = torch.baddbmm(Variable(torch.zeros(max_len, 1).cuda()), encoder_outputs.transpose(0,1), hidden.transpose(0,1)).squeeze(2)

But I got following error and it’s not working.
RuntimeError: expected 3D tensor at /b/wheel/pytorch-src/torch/lib/THC/generic/THCTensorMathBlas.cu:437
If you could figure out how to apply baddmm for my case, could you help me?