Cannot script PyTorch model

[vm3538]

Description

To be able to use a trained DeepAREstimator in Java (Deep Java Library), I wanted to export the model using scripting. However it seems like this is not possible. Are torch scripts supported, or is there a plan to support them?

To Reproduce

Here is the slightly modified script from the GluonTS examples:

import pandas as pd
import torch

from gluonts.dataset.loader import InferenceDataLoader
from gluonts.dataset.pandas import PandasDataset
from gluonts.dataset.split import split
from gluonts.torch import DeepAREstimator
from gluonts.torch.batchify import batchify
from pathlib import Path


df = pd.read_csv(
    'AirPassengers.csv',
    index_col=0,
    parse_dates=True,
)
dataset = PandasDataset(df, target="#Passengers")

training_data, test_gen = split(dataset, offset=-36)
test_data = test_gen.generate_instances(prediction_length=12, windows=3)

model = DeepAREstimator(
    prediction_length=12, freq="M", trainer_kwargs={"max_epochs": 5}
).train(training_data)

model.prediction_net.eval()

with torch.no_grad():
    data_loader = InferenceDataLoader(test_data.input, transform=model.input_transform, batch_size=model.batch_size, stack_fn=lambda data: batchify(data, device=model.device))
    for batch in data_loader:
        inputs = [batch[k] for k in model.input_names]
        trace = model.prediction_net.to_torchscript(method='script', example_inputs=tuple(inputs))
        trace.save('model.pt')
        break

Error message or code output

Traceback (most recent call last):
  File "//trainer/minimal_code.py", line 32, in <module>
    trace = model.prediction_net.to_torchscript(method='script', example_inputs=tuple(inputs))
  File "/usr/local/lib/python3.10/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 1452, in to_torchscript
    torchscript_module = torch.jit.script(self.eval(), **kwargs)
  File "/usr/local/lib/python3.10/site-packages/torch/jit/_script.py", line 1286, in script
    return torch.jit._recursive.create_script_module(
  File "/usr/local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 476, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/usr/local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 488, in create_script_module_impl
    method_stubs = stubs_fn(nn_module)
  File "/usr/local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 757, in infer_methods_to_compile
    stubs.append(make_stub_from_method(nn_module, method))
  File "/usr/local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 69, in make_stub_from_method
    return make_stub(func, method_name)
  File "/usr/local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 54, in make_stub
    ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
  File "/usr/local/lib/python3.10/site-packages/torch/jit/frontend.py", line 293, in get_jit_def
    return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types)
  File "/usr/local/lib/python3.10/site-packages/torch/jit/frontend.py", line 331, in build_def
    param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types)
  File "/usr/local/lib/python3.10/site-packages/torch/jit/frontend.py", line 355, in build_param_list
    raise NotSupportedError(ctx_range, _vararg_kwarg_err)
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/usr/local/lib/python3.10/site-packages/gluonts/torch/model/deepar/lightning_module.py", line 67
    def forward(self, *args, **kwargs):
                              ~~~~~~~ <--- HERE
        return self.model(*args, **kwargs)

If I replace 'script' by 'trace' I won't get any error, but some warnings:

/usr/local/lib/python3.10/site-packages/gluonts/torch/model/deepar/module.py:231: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if future_length > 1:
/usr/local/lib/python3.10/site-packages/gluonts/torch/util.py:137: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert max(indices) <= prior_sequence.shape[dim], (
/usr/local/lib/python3.10/site-packages/torch/jit/_trace.py:1001: TracerWarning: Output nr 1. of the traced function does not match the corresponding output of the Python function. Detailed error:
Tensor-likes are not close!

Mismatched elements: 3599 / 3600 (100.0%)
Greatest absolute difference: 284.1161804199219 at index (0, 19, 7) (up to 1e-05 allowed)
Greatest relative difference: 1.513477324757297 at index (1, 64, 1) (up to 1e-05 allowed)
  _check_trace(

Environment

• Operating system: debian
• Python version: 3.10
• GluonTS version: 0.14.3
• MXNet version: -
• PyTorch version: 1.13.1
• PyTorch Lightning: 2.1.3