@tom Thank you very much, Sir. Yes, I followed what you said. Below is the error message. What am I missing?
TypeError: 'dict' object for attribute 'module_dict' is not a valid constant.
Valid constants are:
  1. a nn.ModuleList
  2. a value of type {bool, float, int, str, NoneType, function, device, layout, dtype}
  3. a list or tuple of (2)

Below is a snippet of combining them
#trace each channel's model separately
channel_tensors = []
module_dict = dict()
for channel_id in channels_list:
    channel = torch.tensor([channel_id], dtype=torch.int64)
    channel_tensors.append(channel)
    traced_script_module = torch.jit.trace(model, (premise, premise_length, hypotheses, hypotheses_length, channel))
    module_dict[channel_id] = traced_script_module

#combine each channel's model together
class MyScriptModule(torch.jit.ScriptModule):
    __constants__ = ['module_dict']

    def __init__(self, module_dict):
        super(MyScriptModule, self).__init__()
        self.module_dict = module_dict

    @torch.jit.script_method
    def forward(self,
                premises,
                premises_lengths,
                hypotheses,
                hypotheses_lengths,
                channel): #channel must be a tensor
        channel_id = channel.item()
        return self.module_dict[channel_id](premise, premise_length, hypotheses, hypotheses_length, channel)

my_script_module = MyScriptModule(module_dict)
my_script_module.save("deploy-trace-multitask.channel_all.pt")