How can I pass keyword arguments to torch.jit.trace_module(model,...), or is it possible?

[i10416]

How can I pass keyword arguments to torch.jit.trace_module(model,...)?

I can pass positional arguments in the same way as the following document and example demonstrate.

• http://djl.ai/docs/pytorch/how_to_convert_your_model_to_torchscript.html
• https://github.com/deepjavalibrary/djl/blob/9d7737cd90f66db9f1d2c30c783c475e65edce28/examples/src/main/java/ai/djl/examples/inference/clip/TextTranslator.java

For example, my model calls generate(inputs,options,...,kwargs*) instead of forward. inputs is a tensor of token_ids obtained from Embedding. There are some keyword arguments such as max_length and temperature, which significantly improve inference performance.

torch.jit.trace_module(model, {"generate": [inputs] } # this works.

However, I couldn't find a way to pass keyword arguments to my model via jit.trace_module.

I tried the following, but it won't work.

torch.jit.trace_module(model, {"method_name":  {"inputs": input_tensor, "keyword0": value0, "keyword1": value1,... }  }

[KexinFeng]

You can take a look at the gpt2 model tracing script in #2509 's PR description. This should solve the problem.