torch2trt icon indicating copy to clipboard operation
torch2trt copied to clipboard

[torch2trt/converters] Add element support to `torch.Tensor.__getitem__` converter

Open chaoz-dev opened this issue 3 years ago • 2 comments

Resolves #768.

This PR adds support for single element arguments to torch.Tensor.__getitem__ converter. We convert the input into a tuple if it is not already a tuple, which can occur when the input argument is a single argument. All other arguments should be handled the same (especially multiple arguments, as they are already tuples).

Single tuple arguments are unpacked correctly still as expected into multiple indexing arguments. Note that tuples combined with ellipsis or colon arguments are still not yet supported; see #755.

Also, add a lot of unit tests to check that operations on the batch dimension are handled correctly.

chaoz-dev avatar Jul 23 '22 04:07 chaoz-dev

Note that this was tested on v0.3.x; the unit tests here fail on v0.4.0. See #769.

chaoz-dev avatar Jul 23 '22 04:07 chaoz-dev

Running python -m torch2trt.test --name getitem:

NUM_TESTS: 39
NUM_SUCCESSFUL_CONVERSION: 39
NUM_FAILED_CONVERSION: 0
NUM_ABOVE_TOLERANCE: 0
NUM_pSNR_TOLERANCE: 0

chaoz-dev avatar Jul 23 '22 04:07 chaoz-dev

NUM_TESTS: 38
NUM_SUCCESSFUL_CONVERSION: 38
NUM_FAILED_CONVERSION: 0
NUM_ABOVE_TOLERANCE: 0
NUM_pSNR_TOLERANCE: 0

chaoz-dev avatar Nov 21 '22 02:11 chaoz-dev