Not able to trace GPT2DoubleHeadsModel

[levindabhi]

System Info

Hi, I'm trying to create trace of GPT2DoubleHeadsModel model but I'm facing issue. Here is my code

from transformers.utils import fx
from transformers import *

gpt2_config = GPT2Config()
model = GPT2DoubleHeadsModel(gpt2_config)
input_names = model.dummy_inputs.keys()
trace = fx.symbolic_trace(model, input_names)

I'm getting below error

File "~/venv/lib/python3.12/site-packages/torch/fx/proxy.py", line 327, in iter
raise TraceError('Proxy object cannot be iterated. This can be '
torch.fx.proxy.TraceError: Proxy object cannot be iterated. This can be attempted when the Proxy is used in a loop or as a *args or **kwargs function argument. See the torch.fx docs on pytorch.org for a more detailed explanation of what types of control flow can be traced, and check out the Proxy docstring for help troubleshooting Proxy iteration errors

Any help, Thanks!

Who can help?

@ArthurZucker @michaelbenayoun

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

from transformers.utils import fx
from transformers import *

gpt2_config = GPT2Config()
model = GPT2DoubleHeadsModel(gpt2_config)
input_names = model.dummy_inputs.keys()
trace = fx.symbolic_trace(model, input_names)

Expected behavior

Expecting it to work without error

[Rocketknight1]

Hi @levindabhi, I don't know if we're strongly supporting torch.fx right now - we're refocusing on making sure everything works with torch.compile instead!