Repository: Project-MONAI/MONAI

Issue:
Cannot convert endoscopic_tool_segmentation model to TensorRT
**Describe the bug**
Cannot convert the endoscopic_tool_segmentation model to TensorRT by running the command below:
```
python -m monai.bundle trt_export --net_id network_def --filepath models/model_trt.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json --precision fp16
```

with error output:
```
/home/liubin/data/tmp/MONAI/monai/networks/utils.py:670: UserWarning: There is no dynamic batch range. The converted model only takes (1, 3, 736,
 480) shape input.
  warnings.warn(f"There is no dynamic batch range. The converted model only takes {input_shape} shape input.")
Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/liubin/data/tmp/MONAI/monai/bundle/__main__.py", line 20, in <module>
    fire.Fire()
  File "/usr/local/lib/python3.8/dist-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/usr/local/lib/python3.8/dist-packages/fire/core.py", line 475, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
  File "/usr/local/lib/python3.8/dist-packages/fire/core.py", line 691, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "/home/liubin/data/tmp/MONAI/monai/bundle/scripts.py", line 1137, in trt_export
    _export(
  File "/home/liubin/data/tmp/MONAI/monai/bundle/scripts.py", line 907, in _export
    net = converter(model=net, **kwargs)
  File "/home/liubin/data/tmp/MONAI/monai/networks/utils.py", line 703, in convert_to_trt
    trt_model = torch_tensorrt.compile(
  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:180] Expected ivalues_maps.count(input) to be true but got false
Could not find torch::jit::Value* 9351 produced from %9351 : Tensor = aten::_convolution(%x_0.2, %self.decoder.blocks.0.convs.conv_0.conv.weight.
3, %self.encoder._conv_stem.bias.33, %11, %11, %11, %9349, %9350, %self.encoder._blocks.0.0.expand_ratio.1, %9349, %9349, %9349, %9349) in loweri
ng graph for mini graph input.

```

**To Reproduce**
Steps to reproduce the behavior:
1. Run `docker pull nvcr.io/nvidian/pytorch:23.03-py3` to get the pytorch docker 23.03 image
2. Start a container with the pytorch 23.03 image
3. Create a local folder **your folder** and go into this folder for test
4. Run `git clone https://github.com/Project-MONAI/MONAI.git`, `cd MONAI;python setup.py develop` to install the latest MONAI
5. Run `pip install fire` to install this optional requirement
6. Back to the **your folder**
7. Download the endoscopic_tool_segmentation bundle by running `python -c "import monai; monai.bundle.download('endoscopic_tool_segmentation', version='0.4.2', bundle_dir='./')"`
8. Run the command line `cd endoscopic_tool_segmentation; python -m monai.bundle trt_export --net_id network_def --filepath models/model_trt.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json --precision fp16`

**Additional Information**
In the UpCat.forward function of [basic_unet](https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/basic_unet.py#L153), if change the type hint of `x_e` to `Optional[torch.Tensor]` and the conditional code `torch.jit.isinstance(x_e, torch.Tensor)` to `x_e is None` like the code in 1.1.0, this question will be fixed.



Discussion:
Hi @wyli @ericspod ,

Do you have any suggestion to modify the typehint of this `x_e` arg?

Thanks in advance.

could you please update the ticket with more information, the steps to reproduce are just unclear. version of torch_tensorrt, monai or docker images, version of bundle these are needed to reproduce.
> could you please update the ticket with more information, the steps to reproduce are just unclear. version of torch_tensorrt, monai or docker images, version of bundle these are needed to reproduce.

Hi @wyli,  I have updated the **To Reproduce** part. Could you please try it and see if you also get this issue?

Resolving patch:
diff --git a/monai/networks/nets/basic_unet.py b/monai/networks/nets/basic_unet.py
--- a/monai/networks/nets/basic_unet.py
+++ b/monai/networks/nets/basic_unet.py
@@ -12,7 +12,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Any
+from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -150,7 +150,7 @@ def __init__(
         self.convs = TwoConv(spatial_dims, cat_chns + up_chns, out_chns, act, norm, bias, dropout)
         self.is_pad = is_pad
 
-    def forward(self, x: torch.Tensor, x_e: Any):
+    def forward(self, x: torch.Tensor, x_e: Optional[torch.Tensor]):
         """
 
         Args:
@@ -159,7 +159,7 @@ def forward(self, x: torch.Tensor, x_e: Any):
         """
         x_0 = self.upsample(x)
 
-        if torch.jit.isinstance(x_e, torch.Tensor):
+        if x_e is not None and torch.jit.isinstance(x_e, torch.Tensor):
             if self.is_pad:
                 # handling spatial shapes due to the 2x maxpooling with odd edge lengths.
                 dimensions = len(x.shape) - 2
diff --git a/monai/networks/nets/dints.py b/monai/networks/nets/dints.py
--- a/monai/networks/nets/dints.py
+++ b/monai/networks/nets/dints.py
@@ -13,6 +13,7 @@
 
 import datetime
 import warnings
+from typing import Optional
 
 import numpy as np
 import torch
@@ -40,7 +41,7 @@
 class CellInterface(torch.nn.Module):
     """interface for torchscriptable Cell"""
 
-    def forward(self, x: torch.Tensor, weight) -> torch.Tensor:  # type: ignore
+    def forward(self, x: torch.Tensor, weight: Optional[torch.Tensor]) -> torch.Tensor:  # type: ignore
         pass
 
 
@@ -170,7 +171,7 @@ def __init__(self, c: int, ops: dict, arch_code_c=None):
             if arch_c > 0:
                 self.ops.append(ops[op_name](c))
 
-    def forward(self, x: torch.Tensor, weight: torch.Tensor | None = None):
+    def forward(self, x: torch.Tensor, weight: Optional[torch.Tensor] = None):
         """
         Args:
             x: input tensor.
@@ -298,7 +299,7 @@ def __init__(
 
         self.op = MixedOp(c, self.OPS, arch_code_c)
 
-    def forward(self, x: torch.Tensor, weight: torch.Tensor | None) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, weight: Optional[torch.Tensor]) -> torch.Tensor:
         """
         Args:
             x: input tensor
diff --git a/monai/networks/nets/dynunet.py b/monai/networks/nets/dynunet.py
--- a/monai/networks/nets/dynunet.py
+++ b/monai/networks/nets/dynunet.py
@@ -9,9 +9,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
+# isort: dont-add-import: from __future__ import annotations
 
-from collections.abc import Sequence
+from typing import List, Optional, Sequence, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -32,7 +32,7 @@ class DynUNetSkipLayer(nn.Module):
     forward passes of the network.
     """
 
-    heads: list[torch.Tensor] | None
+    heads: Optional[List[torch.Tensor]]
 
     def __init__(self, index, downsample, upsample, next_layer, heads=None, super_head=None):
         super().__init__()
@@ -132,13 +132,13 @@ def __init__(
         spatial_dims: int,
         in_channels: int,
         out_channels: int,
-        kernel_size: Sequence[Sequence[int] | int],
-        strides: Sequence[Sequence[int] | int],
-        upsample_kernel_size: Sequence[Sequence[int] | int],
-        filters: Sequence[int] | None = None,
-        dropout: tuple | str | float | None = None,
-        norm_name: tuple | str = ("INSTANCE", {"affine": True}),
-        act_name: tuple | str = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}),
+        kernel_size: Sequence[Union[Sequence[int], int]],
+        strides: Sequence[Union[Sequence[int], int]],
+        upsample_kernel_size: Sequence[Union[Sequence[int], int]],
+        filters: Optional[Sequence[int]] = None,
+        dropout: Optional[Union[Tuple, str, float]] = None,
+        norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}),
+        act_name: Union[Tuple, str] = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}),
         deep_supervision: bool = False,
         deep_supr_num: int = 1,
         res_block: bool = False,
@@ -169,7 +169,7 @@ def __init__(
         self.deep_supervision = deep_supervision
         self.deep_supr_num = deep_supr_num
         # initialize the typed list of supervision head outputs so that Torchscript can recognize what's going on
-        self.heads: list[torch.Tensor] = [torch.rand(1)] * self.deep_supr_num
+        self.heads: List[torch.Tensor] = [torch.rand(1)] * self.deep_supr_num
         if self.deep_supervision:
             self.deep_supervision_heads = self.get_deep_supervision_heads()
             self.check_deep_supr_num()
@@ -305,24 +305,30 @@ def get_output_block(self, idx: int):
     def get_downsamples(self):
         inp, out = self.filters[:-2], self.filters[1:-1]
         strides, kernel_size = self.strides[1:-1], self.kernel_size[1:-1]
-        return self.get_module_list(inp, out, kernel_size, strides, self.conv_block)
+        return self.get_module_list(inp, out, kernel_size, strides, self.conv_block)  # type: ignore
 
     def get_upsamples(self):
         inp, out = self.filters[1:][::-1], self.filters[:-1][::-1]
         strides, kernel_size = self.strides[1:][::-1], self.kernel_size[1:][::-1]
         upsample_kernel_size = self.upsample_kernel_size[::-1]
         return self.get_module_list(
-            inp, out, kernel_size, strides, UnetUpBlock, upsample_kernel_size, trans_bias=self.trans_bias
+            inp,  # type: ignore
+            out,  # type: ignore
+            kernel_size,
+            strides,
+            UnetUpBlock,  # type: ignore
+            upsample_kernel_size,
+            trans_bias=self.trans_bias,
         )
 
     def get_module_list(
         self,
-        in_channels: Sequence[int],
-        out_channels: Sequence[int],
-        kernel_size: Sequence[Sequence[int] | int],
-        strides: Sequence[Sequence[int] | int],
-        conv_block: type[nn.Module],
-        upsample_kernel_size: Sequence[Sequence[int] | int] | None = None,
+        in_channels: List[int],
+        out_channels: List[int],
+        kernel_size: Sequence[Union[Sequence[int], int]],
+        strides: Sequence[Union[Sequence[int], int]],
+        conv_block: nn.Module,
+        upsample_kernel_size: Optional[Sequence[Union[Sequence[int], int]]] = None,
         trans_bias: bool = False,
     ):
         layers = []


Test patch: