I had redefined index so that [4, 1] was not included. My intended point was that rather than assuming each row of a will have 2 elements equal to 1, there is a more general case where any subset of the rows and elements could be chosen. For that setup, I found it more convenient to treat index as a list of (x, y) coordinates i.e.
index = torch.tensor([
[0, 1], # First row, second column
[1, 2], # Second row, third column
# Nothing in the 3rd row
])