Torch.jit.script cannot index a ModuleList with an int returned from range()

Ali-Arch808: <p>Hello everyone,<br>
I’m trying to get a script of a model.<br>
Here is the code(I removed most of the code in <strong>init</strong> function to avoid verbosity):<br>
(spec.matrix is a square matrix that its shape is not predetermined.)</p>
<pre><code class="lang-auto">class Cell(nn.Module):

    def __init__(self, spec, in_channels, out_channels):
        super(Cell, self).__init__()
        self.spec = spec
        self.num_vertices = self.spec.matrix.shape[0]

        # vertex_channels[i] = number of output channels of vertex i
        self.vertex_channels = ComputeVertexChannels(in_channels, out_channels, self.spec.matrix)
        


    def forward(self, x: Any) -&gt; Any:
        tensors = [x]

        out_concat = []

        for t in range(1, self.num_vertices - 1):

            fan_in = []
            for  src in range(1, t):
                if self.spec.matrix[src, t]:
                    fan_in.append(Truncate(tensors[src], torch.tensor(self.vertex_channels[t])))

            if self.spec.matrix[0, t]:
                fan_in.append(self.input_op[t](x))

            # perform operation on node

            vertex_input = sum(fan_in)
            vertex_output = self.vertex_op[t](vertex_input)

            tensors.append(vertex_output)
            if self.spec.matrix[t, self.num_vertices - 1]:
                out_concat.append(tensors[t])

        if not out_concat:
            assert self.spec.matrix[0, self.num_vertices - 1]
            outputs = self.input_op[self.num_vertices - 1](tensors[0])
        else:
            if len(out_concat) == 1:
                outputs = out_concat[0]
            else:
                outputs = torch.cat(out_concat, 1)

            if self.spec.matrix[0, self.num_vertices - 1]:
                outputs += self.input_op[self.num_vertices - 1](tensors[0])


        return outputs

def Truncate(inputs, channels):
    """Slice the inputs to channels if necessary."""
    input_channels = inputs.size()[1]
    if input_channels &lt; channels:
        raise ValueError('input channel &lt; output channels for truncate')
    elif input_channels == channels:
        return inputs  # No truncation necessary
    else:
        # Truncation should only be necessary when channel division leads to
        # vertices with +1 channels. The input vertex should always be projected to
        # the minimum channel count.
        assert input_channels - channels == 1
        return inputs[:, :channels, :, :]
</code></pre>
<p>It returns the following RuntimeError:</p>
<pre><code class="lang-auto">Truncate(Tensor inputs, Tensor channels) -&gt; (Tensor):
Expected a value of type 'Tensor (inferred)' for argument 'inputs' but instead found type 'Any'.
Inferred 'inputs' to be of type 'Tensor' because it was not annotated with an explicit type.:
  File "model.py", line 136
            for  src in range(1, t):
                if self.spec.matrix[src, t]:
                    fan_in.append(Truncate(tensors[src], torch.tensor(self.vertex_channels[t])))
                                  ~~~~~~~~ &lt;--- HERE
    
            if self.spec.matrix[0, t]:
</code></pre>
<p>As error declares, I should annotate the input in tensors. However, doing that will result in <code>RuntimeError:Expected integer literal for index ...</code>.<br>
As I understood from <a href="https://github.com/pytorch/pytorch/issues/47496#issuecomment-789969702" rel="noopener nofollow ugc">this post</a> I tried adding <code>@ torch.jit.interface</code> before the <code>class Cell(nn.Module)</code>. It returned <code>RuntimeError:interface declarations must have a return type annotated.</code> and it underlines all the lines in the <strong>init</strong> function.<br>
Although I can shun the problem by adding a <code>@torch.jit.ignore</code> before the forward function in the Cell Module, it seems a bit of an unreliable solution.<br>
Thank you in advance for any guidance.</p>