Indexing or masking in TRT


I am trying to do simple task which is indexing or masking during using “ONNX_ATEN_FALLBACK” operator at pytorch, I can’t find a way to do it
I tried all of the following solutions :

  • Gather function but it is not supported at TRT yet
  • The normal indexing " [mask]" but it gives an Aten operation during conversion from pytorch to onnx
  • The arange way “x[arange()][idx]” but it gives me the same result at the previous one
  • Maked_select function but it doesn’t work because it is converted to “Nonzero ,Expand” which not supported at TRT

So is there any other solutions without using external written c++ functions?


TensorRT Version: TensorRT
GPU Type: GeForce RTX 2080 Ti/PCIe/SSE2
Nvidia Driver Version: release 460.32.03
CUDA Version: NVIDIA CUDA 11.2.1
CUDNN Version: NVIDIA cuDNN 8.1.0
Operating System + Version: Ubuntu 18.04.3 LTS
Python Version (if applicable): python 3.6 and 3.8
PyTorch Version (if applicable): pytorch 1.7.1 and pytorch 1.8
Baremetal or Container (if container which image + tag): [TensorRT Release 21.03]

Hi @omar.magdy,

Please refer,

Thank you.

Thank you for your suggestion but this one doesn’t support alot of modules used at my project and doesn’t support embedded devices