@tom I also tried another aproach by defining a simple python function to trace, and get errors as below.
    _jit_script_compile(mod, ast, _rcb, get_default_args(fn))
RuntimeError: 
python value of type 'dict' cannot be used as a value:
@torch.jit.script
def loop_in_traced_fn(premise, premise_length, hypotheses, hypotheses_length, channel):
    channel_id = channel.item()
    result = module_dict[channel_id](premise, premise_length, hypotheses, hypotheses_length, channel)
             ~~~~~~~~~~~ <--- HERE
    return result

Below is the code snipeet
#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

@torch.jit.script
def loop_in_traced_fn(premise, premise_length, hypotheses, hypotheses_length, channel):
    channel_id = channel.item()
    result = module_dict[channel_id](premise, premise_length, hypotheses, hypotheses_length, channel)
    return result

channel = channel_tensors[0]
traced = torch.jit.trace(loop_in_traced_fn, (premise, premise_length, hypotheses, hypotheses_length, channel))
traced.save("deploy-trace-multitask.channel_all.pt")