I don’t think this would be possible via direct indexing as the result tensors would have a variable size, wouldn’t they?
In that case you could loop all indices and append the result tensors to a list.