fix: aten.index converter

[zewenli98]

Description

See details in the issue below.

Fixes #2480

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[apbose]

Looks good!

[gs-olive]

Use get for dictionary

[gs-olive]

It seems that the bug in #2480 is only related to the fact that our index converter cannot handle boolean mask indices. The fact that the example has data-dependent shapes does not seem to be the root cause of the initial error.

This is a good temporary fix. Could you also file an issue which specifies that we need to expand index support to boolean masks, since this is a valid usage in Torch and we already support something very similar in aten.where, we just need to add the functionality here. See the sample case below:

>>> import torch
>>> y = torch.rand(2, 2)
>>> y
tensor([[0.3105, 0.9580],
        [0.4824, 0.0796]])
>>> y[torch.Tensor([[True, True], [False, False]]).bool()]
tensor([0.3105, 0.9580])