For all of you that are interested, I think I found a solution, after many trials and errors:
the trick is to set the torch.LongTensor used as index to have the same 0-th dimension value for each dimension of the original tensor.
The code is very simple:
# select first and second row for each batch element
rows = torch.LongTensor( [ [0,1]*batch_size ] ) 
# create an index for each batch element
idx = [
        torch.LongTensor(range(batch_size)).unsqueeze(1), 
        torch.LongTensor(rows) 
      ]
# index the tensor
x[idx] # shape (batch_size, 2, features)

The unsqueeze(1) creates a tensor like this: [[0],[1],[2]...,[batch_size-1]].
By simply changing [0,1] inside the rows tensor it is possible to get different rows for each batch element.