Open
Description
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 🐛 [Bug] Encountered bug when using Torch-TensorRT (We don't have an op for aten::floor_divide but it isn't a special case) #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