I think torch.bmm does what you want. But you’ll need to unsqueeze the query first.
i.e.,
query = query.unsqueeze(2)
result = torch.bmm(sequence,query)
and shape of result would be : batch_size, seq_length,1