antspy:

As I mentioned before, I will consider a “module” everything that is returned by model.modules().


This is still ill-defined. What if a module contains Sequential(linear1, linear2)? All three will show up in model.modules(), but one of them contains the other two. What do you mean by n-th module?



 antspy:

It turns out, the error was that I forgot to de-register the forward hooks. The code below works as intended.
for idx, md in model.named_modules():
if idx == idx_to_stop_execution_to:
md.register_forward_hook(my_function)
try: _ = model(inputs)
except CustomException: pass


This won’t work for a bunch of models, e.g.
def forward(self, x):
  for _ in range(100):
    x = self.fc(x)
  return x

def forward(self, x):
  if x.data[0] < 1:
    return self.fc1(x)
  else:
    return self.fc2(x)
  return x

def forward(self, x):
  z = self.subnet(Variable(torch.randn(1, 2))
  return z + x

No where guarantees that named_modules or modules returns in certain order. It doesn’t even work for sequential modules.
If you want to limit youself to nn.Sequential, fine, there are a number of ways to do this. But if you want to do it for general modules. I don’t see how any of the proposals you made could work. You should think about what exactly is a “module”,  and with dynamic python flow control, how to get the n-th actually executed “module”.