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!

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.