vmap over getitem indexing raises a RunTime error

[ordabayevy]

🐛 Describe the bug

Calling torch.vmap over indexing raises a runtime error:

import torch


def index(x, y):
    return x[y]

batched_index = torch.vmap(index, (None, 0))

x = torch.arange(3)

# this works as expected
# tensor([[0, 2],
#         [1, 2]])
y1 = torch.tensor([[0, 2], [1, 2]])
result1 = batched_index(x, y1)
print(result1)

# this raises RuntimeError: vmap: It looks like you're calling .item() on a Tensor. We don't support vmap over calling .item()
# on a Tensor, please try to rewrite what you're doing with other operations. If error is occurring somewhere inside PyTorch 
# internals, please file a bug report.
y2 = torch.tensor([0, 2])
result2 = batched_index(x, y2)
print(result2)

Expected result2 to return torch.tensor([0, 2])

This is related to #115347 but not exactly the same.

Versions

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] optree==0.11.0
[pip3] torch==2.4.0a0+git9c4fc5f
[conda] magma-cuda121             2.6.1                         1    pytorch
[conda] magma-cuda124             2.6.1                         1    pytorch
[conda] mkl-include               2024.1.0              intel_691    intel
[conda] mkl-static                2024.1.0              intel_691    intel
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] torch                     2.4.0a0+git9c4fc5f           dev_0    <develop>

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99