akankshakshurai/cuda-error-resolution-analysis huggingface dataset:
Torch.split() in torch_tensorrt --> Compile Error

anass_017: <p>i have got a problem in the TensorRT  when using torch.split() on one of my Layers</p>
<p>The Class</p>
<pre><code class="lang-auto">class Unpack(nn.Module):
    def __init__(self,split_size=3):
        super(Unpack,self).__init__()
        self.split_size = split_size

    def forward(self,input):
        input_r,input_i =torch.split(input,split_size_or_sections=self.split_size,dim=0)
        return input_r,input_i
</code></pre>
<p>The Code part :</p>
<pre><code class="lang-auto">trt_model_fp16 = torch_tensorrt.compile(model,
                                            inputs = [(torch_tensorrt.Input((1,)+X_cat.shape , dtype=data_type_16))],
                                            enabled_precisions = data_type_16, # Run with FP16
                                            workspace_size = 1 &lt;&lt; 22)
</code></pre>
<p>The Error (I think the error is on the TorchScript conversion)</p>
<pre><code class="lang-auto"> if module_type == _ModuleType.nn:
    120         logging.log(
    121             logging.Level.Info,
    122             "Module was provided as a torch.nn.Module, trying to script the module with torch.jit.script. In the event of a failure please preconvert your module to TorchScript",
    123         )
--&gt; 124         ts_mod = torch.jit.script(module)
    125     return torch_tensorrt.ts.compile(
    126         ts_mod, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs
    127     )
    128 elif target_ir == _IRType.fx:
...
        #print(input.shape)
        input_r,input_i =torch.split(input,split_size_or_sections=self.split_size,dim=1)
                         ~~~~~~~~~~~ &lt;--- HERE
        return input_r,input_i

</code></pre>

ptrblck: <p>CC <a class="mention" href="/u/narendasan">@narendasan</a> as the expert on TorchTRT questions.</p>

narendasan: <p><a class="mention" href="/u/anass_017">@anass_017</a> Can you file an issue in the Github repo with the full error? I suspect this might be an issue in the graph partitioning phase and not necessarily the torchscript conversion but the full log would make that clear.</p>