torch.arange(A.size(0)) will return tensor([0, 1]) in this case, so the indexing will work as:
out[0] = A[0, index[0]]
out[1] = A[1, index[1]]

and will thus yield the desired results. On the other hand using A[:, index] will not work (you might be more familiar with this indexing operation), as it would use index on “all” values on dim0, so it’ll return:
A[:, index]
> tensor([[[ 1,  2],
           [ 5,  6]],

          [[ 9, 10],
           [13, 14]]])

As you can see, index (containing [0, 2]) was now applied to all values in dim0 as:
A[0, index]
A[1, index]

thus returning more values.
EDIT: this numpy doc 1 might be also helpful in case my explanation is confusing.