Get_graph_node_names() with huggingface transformer models

Hey, I’m trying to use TorchVision’s get_graph_node_names() to fetch the names of nodes available for feature extraction in the Huggingface’s gpt2-small model using the following code:

from torchvision.models.feature_extraction import get_graph_node_names
nodes, _ = get_graph_node_names(gpt_small)
print(nodes)

However, I’m getting the following error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [14], in <cell line: 2>()
      1 from torchvision.models.feature_extraction import get_graph_node_names
----> 2 nodes, _ = get_graph_node_names(gpt_small)
      3 print(nodes)

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/torchvision/models/feature_extraction.py:253, in get_graph_node_names(model, tracer_kwargs, suppress_diff_warning)
    251 is_training = model.training
    252 train_tracer = NodePathTracer(**tracer_kwargs)
--> 253 train_tracer.trace(model.train())
    254 eval_tracer = NodePathTracer(**tracer_kwargs)
    255 eval_tracer.trace(model.eval())

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py:587, in Tracer.trace(self, root, concrete_args)
    585     for module in self._autowrap_search:
    586         _autowrap_check(patcher, module.__dict__, self._autowrap_function_ids)
--> 587     self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
    588                      type_expr=fn.__annotations__.get('return', None))
    590 self.submodule_paths = None
    592 return self.graph

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/transformers/models/gpt2/modeling_gpt2.py:1047, in GPT2LMHeadModel.forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1039 r"""
   1040 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
   1041     Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
   1042     `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
   1043     are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
   1044 """
   1045 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-> 1047 transformer_outputs = self.transformer(
   1048     input_ids,
   1049     past_key_values=past_key_values,
   1050     attention_mask=attention_mask,
   1051     token_type_ids=token_type_ids,
   1052     position_ids=position_ids,
   1053     head_mask=head_mask,
   1054     inputs_embeds=inputs_embeds,
   1055     encoder_hidden_states=encoder_hidden_states,
   1056     encoder_attention_mask=encoder_attention_mask,
   1057     use_cache=use_cache,
   1058     output_attentions=output_attentions,
   1059     output_hidden_states=output_hidden_states,
   1060     return_dict=return_dict,
   1061 )
   1062 hidden_states = transformer_outputs[0]
   1064 # Set device for model parallelism

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py:577, in Tracer.trace.<locals>.module_call_wrapper(mod, *args, **kwargs)
    573     return _orig_module_call(mod, *args, **kwargs)
    575 _autowrap_check(patcher, getattr(getattr(mod, "forward", mod), "__globals__", {}),
    576                 self._autowrap_function_ids)
--> 577 return self.call_module(mod, forward, args, kwargs)

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/torchvision/models/feature_extraction.py:84, in NodePathTracer.call_module(self, m, forward, args, kwargs)
     82 self.current_module_qualname = module_qualname
     83 if not self.is_leaf_module(m, module_qualname):
---> 84     out = forward(*args, **kwargs)
     85     return out
     86 return self.create_proxy("call_module", module_qualname, args, kwargs)

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py:573, in Tracer.trace.<locals>.module_call_wrapper.<locals>.forward(*args, **kwargs)
    572 def forward(*args, **kwargs):
--> 573     return _orig_module_call(mod, *args, **kwargs)

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/miniconda3/envs/biofeedback/lib/python3.10/site-packages/transformers/models/gpt2/modeling_gpt2.py:768, in GPT2Model.forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions, output_hidden_states, return_dict)
    765 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    767 if input_ids is not None and inputs_embeds is not None:
--> 768     raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
    769 elif input_ids is not None:
    770     input_shape = input_ids.size()

ValueError: You cannot specify both input_ids and inputs_embeds at the same time

I don’t understand why both input_ids and inputs_embeds are getting passed to the gpt2-small during the symbolic tracing. Any help would be greatly appreciated.

Pytorch version: 1.12.0

Thanks!