can get_submodule be called within a ScriptFunction ?

[wellhowtosay]

🚀 The feature, motivation and pitch

I'm writing a multi domain classifier where can first predict domains of a sentence, and then choose the corresponding sequence classifer to run. It can certainly save FPOP in inference stage.

class MultiDomainClassifier(nn.Module):
    def __init__(self, classifiers: nn.ModuleDict, num_domain):
        super().__init__()
        self.classifiers = nn.ModuleList(classifiers.children())
        self.num_domain = num_domain

    def forward(self, inputs: Tensor, output_domains: Tensor) -> list[Tensor]:
        return [self.classifiers.get_submodule(str(idx)).forward(inputs) for idx in output_domains]

But I got this:

RuntimeError: 
Unknown type name 'Module':
  File "/home/luka.bao/anaconda3/envs/cpu/lib/python3.9/site-packages/torch/nn/modules/module.py", line 617
    def get_submodule(self, target: str) -> "Module":
                                            ~~~~~~~ <--- HERE
        """
        Returns the submodule given by ``target`` if it exists,
'ModuleList.get_submodule' is being compiled since it was called from 'MultiDomainClassifier.forward'
  File "/home/luka.bao/PycharmProjects/multitask-nlu-torch/torch_impl/output/classifier.py", line 109
    def forward(self, inputs: Tensor, output_domains: Tensor) -> list[Tensor]:
        return [self.classifiers.get_submodule(str(idx)).forward(inputs) for idx in output_domains]
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

Alternatives

No response

Additional context

No response

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel

[wellhowtosay]

Unknown type name 'Module':

thanks for the comment, and I changed all list[Tensor]-> List[Tensor] but I still get the same RuntimeError

RuntimeError: 
Unknown type name 'Module':
  File "/home/luka.bao/anaconda3/envs/cpu/lib/python3.9/site-packages/torch/nn/modules/module.py", line 617
    def get_submodule(self, target: str) -> "Module":
                                            ~~~~~~~ <--- HERE
        """
        Returns the submodule given by ``target`` if it exists,
'ModuleList.get_submodule' is being compiled since it was called from 'MultiDomainClassifier.forward'
  File "/home/luka.bao/PycharmProjects/multitask-nlu-torch/torch_impl/output/classifier.py", line 109
    def forward(self, inputs: Tensor, output_domains: Tensor) -> list[Tensor]:
        return [self.classifiers.get_submodule(str(idx)).forward(inputs) for idx in output_domains]
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

it says Module::get_submodule returns a Unknown type name 'Module'