TensorRT: 🐛 [Bug] Cannot compile SwinIR model (shape_analysis.cpp: Expected ivalues_maps.count(input) to be true but got false)

Bug Description

Cannot compile the SwinIR model.

Error message:

Traceback (most recent call last):
  File "main.py", line 61, in <module>
    compile_tensorrt_model(torch.float)
  File "main.py", line 56, in compile_tensorrt_model
    compiled_model = torch_tensorrt.compile(traced_model, inputs=inputs, enabled_precisions=enabled_precisions,
  File "/usr/local/lib/python3.8/dist-packages/torch_tensorrt/_compile.py", line 125, in compile
    return torch_tensorrt.ts.compile(
  File "/usr/local/lib/python3.8/dist-packages/torch_tensorrt/ts/_compiler.py", line 136, in compile
    compiled_cpp_mod = _C.compile_graph(module._c, _parse_compile_spec(spec))
RuntimeError: [Error thrown at core/partitioning/shape_analysis.cpp:167] Expected ivalues_maps.count(input) to be true but got false
Could not find torch::jit::Value* 71852 produced from %71852 : Tensor = aten::add(%71851, %71850, %71848) in lowering graph for mini graph input.

To Reproduce

The original code is not properly typed, so I modified it a bit. Repo: https://github.com/arition/SwinIR-TensorRT

What I changed compared to original code:

• Add proper typing
• Disable use_checkpoint
• Disable other variants except real-world SR
• Replace modulo operator to custom function according to https://github.com/pytorch/TensorRT/issues/1305

To reproduce, just download pretrained weight (link in code) and run main.py.

Expected behavior

The model compiles without problems.

Environment

I use PyTorch container 23.01-py3 on NGC: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch

• CPU Architecture: x64
• OS (e.g., Linux): Linux
• GPU models and configuration: RTX 4090