How did you script the generate method?
I use torch.jit.script(model.generate) and got and error:
NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults: ...