minimaxir/aitextgen

Torchscript Text Generation Support

minimaxir opened this issue · 1 comments

Although we can export and load Torchscript models, the forward() function of a Torchscript model must match what is generated.

The traced TS model has a forward() w/ 2 parameters, but the raw forward() function of GPT2LMModel has 9 parameters.

Solution may be to use an extended Model class just for TorchScript, with an override forward() function w/ 2 parameters.

Closing and unpinning since ONNX support will supersede TorchScript.