Hello
I’m using 1.5.0 version of PyTorch.
I’d like to skip i numbers of layers in a Torchscript model.
Below code is my trial. The method ‘skip’ takes input tensor x and skip point indicator i.
import torch
import torch.nn as nn

class Model(torch.nn.Module):
        def __init__(self):
            super(Model, self).__init__()
            self.fc1 = nn.Linear(20,100)
            self.fc2 = nn.Linear(100, 100)
            self.fc3 = nn.Linear(100, 1)
            self.k = nn.ModuleList()
            self.k.append(self.fc1)
            self.k.append(self.fc2)
            self.k.append(self.fc3)

        @torch.jit.export
        def skip(self, x, i):
            for l in range(int(i), 3):
                x = self.k[l].forward(x)
            return x

        def forward(self, x):
            x = self.fc1(x)
            x = self.fc2(x)
            x = self.fc3(x)
            return x

model = Model().cuda()
x = torch.randn(20).cuda()
x2 = torch.randn(100).cuda() # Assume that x2 is the result of self.fc1 which is computed in advance.
idx = torch.tensor(1).cuda()
print(model.skip(x2,idx)) # This works

traced = torch.jit.trace(model, x) # Script
print(traced.skip(x2,idx)) # This does not works. Error occurs.

The error message is printed as below.
Traceback (most recent call last):
  File "C:/Users/user/PycharmProjects/coin-modelgen/simple.py", line 33, in <module>
    traced = torch.jit.trace(model, x) # Script
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\__init__.py", line 875, in trace
    check_tolerance, _force_outplace, _module_class)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\__init__.py", line 1021, in trace_module
    module = make_module(mod, _module_class, _compilation_unit)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\__init__.py", line 716, in make_module
    return torch.jit._recursive.create_script_module(mod, make_stubs_from_exported_methods, share_types=False)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\_recursive.py", line 305, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\_recursive.py", line 361, in create_script_module_impl
    create_methods_from_stubs(concrete_type, stubs)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\_recursive.py", line 279, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError: 
Expected integer literal for index:
  File "C:/Users/user/PycharmProjects/coin-modelgen/simple.py", line 18
        def skip(self, x, i):
            for l in range(int(i), 3):
                x = self.k[l].forward(x)
                    ~~~~~~~~ <--- HERE
            return x

Could someone can share any idea (or alternative approach) to solve this problem?
Moreover, I want to know that such dynamically skipping is also possible or not in c++ implementation using libtorch.
According to torch.jit.script cannot index a ModuleList with an int returned from range() · Issue #47496 · pytorch/pytorch · GitHub, I identified that the i must be selected in compile time not runtime. Thus…, is this meaning that there no way to dynamically skip the front part of the Torchscript model?
Or, is such naive below code a unique way?
import torch
import torch.nn as nn

class Model(torch.nn.Module):
        def __init__(self):
            super(Model, self).__init__()
            self.fc1 = nn.Linear(20,100)
            self.fc2 = nn.Linear(100, 100)
            self.fc3 = nn.Linear(100, 1)
            self.k = nn.ModuleList()
            self.k.append(self.fc1)
            self.k.append(self.fc2)
            self.k.append(self.fc3)
            # self.k = nn.Sequential(self.fc1,self.fc2,self.fc3 )

        def from1(self, x):
            for index, v in enumerate(self.k[1:]):
                x = v.forward(x)
            return x

        def from2(self, x):
            for index, v in enumerate(self.k[2:]):
                x = v.forward(x)
            return x

        @torch.jit.export
        def skip(self, x, i:int):
            if i == 1:
                x = self.from1(x)
            elif i == 2:
                x = self.from2(x)
            return x

        def forward(self, x):
            x = self.fc1(x)
            x = self.fc2(x)
            x = self.fc3(x)
            return x

model = Model().cuda()
m = torch.jit.script(Model()).cuda()

x = torch.randn(20).cuda()
x2 = torch.randn(100).cuda() # Assume that x is the result of self.fc1 which is computed in advance.
idx = 1
print(model.skip(x2,idx)) # This works
print(m.skip(x2,idx)) # This works

tensor([-0.3310], device='cuda:0', grad_fn=<AddBackward0>)
tensor([-0.2432], device='cuda:0', grad_fn=<AddBackward0>)

If this way is the unique way, why the two result are different?
Thanks a lot.