Hello @KFrank !
Thank you for your answer. I try to explain a little better what my problem is because maybe I haven’t been very detailed.
At the beginning of my code, I have a tensor with shape x = torch.Size([50, 61, 140]) (batch_size, seq_len, embedding_dim) and a ndarray x_len = (50,). These two vectors I give them as input to a bilstm whose code is the following (I’ll save you all the forward function, I don’t think it’s important):
x = nn.utils.rnn.pack_padded_sequence(x, x_len, batch_first=True)
out1 = self.lstm1(x)
x, lengths = nn.utils.rnn.pad_packed_sequence(out1, batch_first=True)

After this code I have two tensors with shape x = ([50, 61, 800]) and x_len = ([50]).
I give these two vectors as input to a module to calculate the attention in the following way (x is the parameter corresponding to inputs and x_len is the one corresponding to lengths):
def forward(self, inputs, lengths):
    batch_size, max_len = inputs.size()[:2]  # batch_size = 50, max_len = 61

    # matrix mult
    # apply attention layer
    weights = torch.bmm(inputs,
                        self.att_weights  # (1, hidden_size)
                        .permute(1, 0)  # (hidden_size, 1)
                        .unsqueeze(0)  # (1, hidden_size, 1)
                        .repeat(batch_size, 1, 1)  # (batch_size, hidden_size, 1)
                        )

    # weights.shape = torch.Size([50, 61, 1])
    attentions = torch.softmax(F.relu(weights.squeeze()), dim=-1)   # torch.Size([50, 61])

    # create mask based on the sentence lengths
    mask = torch.ones(attentions.size(), requires_grad=True).cuda()  # torch.Size([50, 61])
    for i, l in enumerate(lengths):  # skip the first sentence
        if l < max_len:
            mask[i, l:] = 0

    # apply mask and renormalize attention scores (weights)
    masked = attentions * mask
    _sums = masked.sum(-1).unsqueeze(-1)  # sums per row    # torch.Size([50, 1])

    attentions = masked.div(_sums)   # torch.Size([50, 61])

    # apply attention weights
    weighted = torch.mul(inputs, attentions.unsqueeze(-1).expand_as(inputs))   # torch.Size([50, 61, 800])

    # get the final fixed vector representations of the sentences
    representations = weighted.sum(1).squeeze()   # ([50, 800])

    return representations, attentions   # representations =  ([50, 800]) , attentions = ([50, 61])

What I would like to obtain as the output of this forward function is a tensor with shape ([50, 61, 800]) instead of two tensors made in that way.
Do you think there is a way to do it without “affecting” the attention calculation?
Thank you so much,
best regards.