Unable to convert Huggingface model to torchscript

Beginners
1

Hi Team,

I’m working on Huggingface Tapas model, as it is working, I’m trying to convert Huggingface Tapas model to torchscript model, in order to deploy this model in Nvidia Triton Server.

from transformers import TapasTokenizer, TapasForQuestionAnswering
import pandas as pd

model_name = "google/tapas-base-finetuned-wtq"
model = TapasForQuestionAnswering.from_pretrained(model_name, torchscript=True)
tokenizer = TapasTokenizer.from_pretrained(model_name)
data = {"Actors": ["Brad Pitt", "Leonardo Di Caprio", "George Clooney"], "Number of movies": ["87", "53", "69"]}
queries = "What is the name of the first actor?"
table = pd.DataFrame.from_dict(data)
inputs = tokenizer(table=table, queries=queries, padding="max_length", return_tensors="pt")
outputs = model(**inputs,return_dict=True)
predicted_answer_coordinates, predicted_aggregation_indices = tokenizer.convert_logits_to_predictions(
    inputs, outputs.logits.detach(), outputs.logits_aggregation.detach()
)

# Creating the trace
import torch
traced_model = torch.jit.trace(model, [inputs['input_ids'], inputs['token_type_ids'], inputs['attention_mask']])
torch.jit.save(traced_model, "traced_bert.pt")

when I try to trace the model, its throwing below error which indexerror, could you please help me here? I don’t know whether I’m doing a mistake or not. Please correct me if I’m doing any mistake.

IndexError                                Traceback (most recent call last)
Cell In[10], line 3
      1 # Creating the trace
      2 import torch
----> 3 traced_model = torch.jit.trace(model, [inputs['input_ids'], inputs['token_type_ids'], inputs['attention_mask']])
      4 torch.jit.save(traced_model, "traced_bert.pt")

File ~/miniconda3/envs/scrubai/lib/python3.9/site-packages/torch/jit/_trace.py:735, in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit)
    732     return func
    734 if isinstance(func, torch.nn.Module):
--> 735     return trace_module(
    736         func,
    737         {"forward": example_inputs},
    738         None,
    739         check_trace,
    740         wrap_check_inputs(check_inputs),
    741         check_tolerance,
    742         strict,
    743         _force_outplace,
    744         _module_class,
    745     )
    747 if (
    748     hasattr(func, "__self__")
    749     and isinstance(func.__self__, torch.nn.Module)
...
--> 312     col_index = IndexMap(token_type_ids[:, :, 1], self.config.type_vocab_sizes[1], batch_dims=1)
    313     # shape (batch_size, seq_len)
    314     row_index = IndexMap(token_type_ids[:, :, 2], self.config.type_vocab_sizes[2], batch_dims=1)

IndexError: too many indices for tensor of dimension 2
2

can anyone help me here @sgugger @thomwolf @lysandre @nielsr