Error with jit script

i am trying to save a model after training using JIT, but it throws me this error:

NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parallel/", line 147
    def forward(self, *inputs, **kwargs):
                                ~~~~~~~ <--- HERE
        with torch.autograd.profiler.record_function("DataParallel.forward"):
            if not self.device_ids:

here’s the model,

class custom_model(nn.Module):
    def __init__(self):
        self.base = base_model
        self.mean_max_pool = MeanMaxPool(hidden_size)
        self.dropout = nn.Dropout(0.1)
        self.fc = nn.Linear(hidden_size*3,inter_size[0])
        self.out = nn.Linear(inter_size[0],n_classes)

    def forward(self, ids,mask):
        x = self.base(ids, mask)[0]
        x = self.mean_max_pool(x,mask)
        x= self.fc(x)
        return out

model = custom_model()
model= nn.DataParallel(model,device_ids=[0, 1])


Could you try to script the model before wrapping it into DataParallel?

Thank you for replying, it throws a different error:

UnsupportedNodeError: GeneratorExp aren't supported:
  File "/opt/conda/lib/python3.7/site-packages/transformers/", line 1538
        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
                   ~ <--- HERE

It seems the error is raised by the transformers package and I don’t know what known scripting limitations there are, so you might need to ask in the HuggingFace discussion board.

1 Like