🐛 [Bug] Can't convert SSDLite320 MobilenetV3, Unsupported operator: aten::index.Tensor(Tensor self, Tensor?[] indices) -> (Tensor)

Bug Description

I trained a ssdlite320_320 mobilenetv3 large with Widerface datasets for face detection task.

Here is what I received when running the torch_tensorrt.compile():

(capstone) jetson@jetson-desktop:~/FaceRecognitionSystem/jetson/backend/python$ python test.py
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py:1102: UserWarning: torch.meshgrid: in an upcoming re
lease, it will be required to pass the indexing argument. (Triggered internally at /media/nvidia/NVME/pytorch/pytorch-v1.10.0/aten/src/ATen/native/Te
nsorShape.cpp:2157.)
return forward_call(*input, **kwargs)
Wrapper works successfully!
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/jit/_trace.py:965: TracerWarning: Encountering a list at the output of th
e tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider u
sing a constant container instead (e.g. for list, use a tuple instead. for dict, use a NamedTuple instead). If you absolutely need this and kn
ow the side effects, pass strict=False to trace() to allow this behavior.
argument_names,
Tracing successful!
WARNING: [Torch-TensorRT] - For input x, found user specified input dtype as Float16, however when inspecting the graph, the input type expected was i
nferred to be Float
The compiler is going to use the user setting Float16
This conflict may cause an error at runtime due to partial compilation being enabled and therefore
compatibility with PyTorch's data type convention is required.
If you do indeed see errors at runtime either:

• Remove the dtype spec for x
• Disable partial compilation by setting require_full_compilation to True

ERROR: [Torch-TensorRT] - Unsupported operator: aten::index.Tensor(Tensor self, Tensor?[] indices) -> (Tensor)
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torchvision-0.11.1-py3.6-linux-aarch64.egg/torchvision/models/detection/ssd.py(
406): postprocess_detections
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torchvision-0.11.1-py3.6-linux-aarch64.egg/torchvision/models/detection/ssd.py(
354): forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1090): _slow_forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1102): _call_impl
convert_to_trt.py(36): forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1090): _slow_forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1102): _call_impl
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/jit/_trace.py(965): trace_module
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/jit/_trace.py(750): trace
convert_to_trt.py(54):
Serialized File "code/torch/torchvision/models/detection/ssd.py", line 110
keep0 = torch.slice(keep, 0, 0, 300)
_80 = annotate(List[Optional[Tensor]], [keep0])
boxes2 = torch.index(image_boxes, _80)
~~~~~~~~~~~ <--- HERE
_81 = annotate(List[Optional[Tensor]], [keep0])
_82 = torch.index(image_scores, _81)
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torchvision-0.11.1-py3.6-linux-aarch64.egg/torchvision/models/detection/ssd.py(
408): postprocess_detections
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torchvision-0.11.1-py3.6-linux-aarch64.egg/torchvision/models/detection/ssd.py(
354): forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1090): _slow_forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1102): _call_impl
convert_to_trt.py(36): forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1090): _slow_forward
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/nn/modules/module.py(1102): _call_impl
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/jit/_trace.py(965): trace_module
/home/jetson/miniforge-pypy3/envs/capstone/lib/python3.6/site-packages/torch/jit/_trace.py(750): trace
convert_to_trt.py(54):
Serialized File "code/torch/torchvision/models/detection/ssd.py", line 114
_82 = torch.index(image_scores, _81)
_83 = annotate(List[Optional[Tensor]], [keep0])
_84 = torch.index(image_labels, _83)
~~~~~~~~~~~ <--- HERE
_85 = torch.to(torch.detach(s), torch.device("cuda:0"), 6, False, True)
_86 = torch.detach(_85)

To Reproduce

Steps to reproduce the behavior:

1. Require jetson nano 4GB, torch_tensorrt 1.0.0, torch 1.10.0 and torchvision 0.11.1.
2. Use torchvision implementation of ssdlite320 and convert to tensorrt:

import torch
import torchvision
from torchvision import models
from torch2trt import torch2trt

class FaceDetectionModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = models.detection.ssdlite320_mobilenet_v3_large(num_classes=2)

    def forward(self, images, targets=None):
        if self.training:
            outputs = self.model(images, targets)

            return outputs['bbox_regression'], outputs['classification']
        else:
            outputs = self.model(images) 

            boxes = [out["boxes"] for out in outputs]
            scores = [out["scores"] for out in outputs]
            labels = [out["labels"] for out in outputs]

            return boxes, scores, labels
model = FaceDetectionModel()

checkpoint = torch.load("./saved_model/face_detection3_epoch200_loss0.1802.pth")
new_state_dict = {k.replace("module.", ""): v for k, v in checkpoint.items()}
model.load_state_dict(new_state_dict)
model.eval().to("cuda")

dummy_input = torch.randn(1, 3, 320, 320).to("cuda")

# Tracing
with torch.no_grad():
    traced_model = torch.jit.trace(model, (dummy_input,))

# Save
traced_model.save("saved_model/face_detection3_epoch200_loss0.1802_traced.ts")
torch.save(model, "./saved_model/face_detection3_epoch200_loss0.1802.pt")

model = torch.jit.load('saved_model/face_detection3_epoch200_loss0.1802_traced.ts')
model.eval().to("cuda")

inputs=[
    torch_tensorrt.Input((1, 3, 320, 320), dtype=torch.float16)
]

enabled_precisions={torch.float16}  # Use FP16 for faster inference
# Compile with TensorRT
trt_model = torch_tensorrt.compile(
    model,
    inputs = inputs,
    enabled_precisions=enabled_precisions,
    # require_full_compilation = False # This is enable by default
)

input_data = torch.randn(1, 3, 320, 320).to("cuda").half()

result = trt_model(input_data)

# Save the TensorRT optimized model
torch.jit.save(trt_model, "saved_model/trt_optimized_model.ts")

3. Run the script and return the error.

I expected ssdlite320 mobilenetv3 should be supported to convert to tensorrt

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Torch-TensorRT Version (e.g. 1.0.0): 1.0.0
• PyTorch Version (e.g. 1.0): 1.10.0
• CPU Architecture: ARM
• OS (e.g., Linux): Ubuntu 18.04
• How you installed PyTorch (conda, pip, libtorch, source): source
• Python version: 3.6.15
• CUDA version: 10.2