weidler:

by selectively picking along the third domain based on the position in the second domain.


I don’t quite understand this use case, but based on your loop approach, you could use arange to index the tensor via:
original_tensor[:, torch.arange(original_tensor.size(1)), torch.arange(original_tensor.size(2))]

which would result in the same output as output_tensor.