"index with mask" failure of TensorRT 8.6 when running trtexec on GPU A4500
Description
I tried to get tensor with index,codes is
"""
class SimpleModel(nn.Module):
def forward(self, x, mask):
out = x[mask]
return out
"""
The onnx graph is
A bug from onnx to trt
Environment
docker is :nvcr.io/nvidia/pytorch:23.04-py3
Relevant Files
Model link: test.tar.gz
Steps To Reproduce
trtexec --onnx=simp_model_sim.onnx --saveEngine=tmp.trt
Commands or scripts:
Have you tried the latest release?: My environment does not support using more advanced versions.
Can this model run on other frameworks? For example run ONNX model with ONNXRuntime (polygraphy run <model.onnx> --onnxrt):
onnxrt is PASSED.
I think it related to data-depended shape.
@lix19937 My tensorrt version is 8.6.1, I think it supported data-depended shape, like non_zero, where
You can try the trt latest version.
Hello there's a sample for data-dependent shape plugins here. You can also override IOutputAllocator for this purpose. Closing for now but feel free to reopen if you still see the issue for trt >=10.0!