Control flow issue with symbolic_trace when using inputs_embeds in LlamaForCausalLM
Author: Hongjie1Chu
Pull request: no

### System Info


- `transformers` version: 4.41.2
- Platform: Linux-5.15.0-88-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.23.0
- Safetensors version: 0.4.3
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>


### Who can help?

@fxmarty and @michaelbenayoun @younesbelkada

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

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

### Reproduction

```
from transformers import LlamaForCausalLM,LlamaConfig
import torch
from transformers.utils.fx import symbolic_trace


cfg = LlamaConfig()
cfg.num_hidden_layers = 1
model = LlamaForCausalLM(cfg)
print(model)
traced = symbolic_trace(model,input_names=['input_ids'])
traced = symbolic_trace(model,input_names=['inputs_embeds'])

```

### Expected behavior

I would like to know why using input_ids allows for normal tracing, and there are no errors when executing '        if query_states.device.type == "cuda" and causal_mask is not None: ', but switching to inputs_embeds causes an issue at this point?

the error is:

---------------------------------------------------------------------------
TraceError                                Traceback (most recent call last)
Cell In[7], [line 1](vscode-notebook-cell:?execution_count=7&line=1)
----> [1](vscode-notebook-cell:?execution_count=7&line=1) traced = symbolic_trace(model,input_names=['inputs_embeds'])

File ~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1483, in symbolic_trace(model, input_names, disable_check, tracer_cls)
   [1481](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1481) # Tracing.
   [1482](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1482) tracer = tracer_cls()
-> [1483](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1483) traced_graph = tracer.trace(model, concrete_args=concrete_args)
   [1484](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1484) traced = torch.fx.GraphModule(model, traced_graph)
   [1486](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1486) traced.config = model.config

File ~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1306, in HFTracer.trace(self, root, concrete_args, dummy_inputs, complete_concrete_args_with_inputs_not_in_dummy_inputs)
   [1304](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1304) with self.patch_for_tracing(root):
   [1305](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1305)     try:
-> [1306](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1306)         self.graph = super().trace(root, concrete_args=concrete_args)
   [1307](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1307)     finally:
   [1308](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1308)         _CURRENT_TRACER = None

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:817, in Tracer.trace(self, root, concrete_args)
    [810](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:810)         for module in self._autowrap_search:
    [811](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:811)             _autowrap_check(
    [812](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:812)                 patcher, module.__dict__, self._autowrap_function_ids
    [813](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:813)             )
    [814](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:814)         self.create_node(
    [815](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:815)             "output",
    [816](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:816)             "output",
--> [817](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:817)             (self.create_arg(fn(*args)),),
    [818](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:818)             {},
    [819](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:819)             type_expr=fn.__annotations__.get("return", None),
    [820](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:820)         )
    [822](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:822)     self.submodule_paths = None
    [823](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:823) finally:

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1168, in LlamaForCausalLM.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
   [1165](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1165) return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   [1167](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1167) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> [1168](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1168) outputs = self.model(
   [1169](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1169)     input_ids=input_ids,
   [1170](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1170)     attention_mask=attention_mask,
   [1171](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1171)     position_ids=position_ids,
   [1172](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1172)     past_key_values=past_key_values,
   [1173](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1173)     inputs_embeds=inputs_embeds,
   [1174](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1174)     use_cache=use_cache,
   [1175](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1175)     output_attentions=output_attentions,
   [1176](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1176)     output_hidden_states=output_hidden_states,
   [1177](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1177)     return_dict=return_dict,
   [1178](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1178)     cache_position=cache_position,
   [1179](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1179) )
   [1181](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1181) hidden_states = outputs[0]
   [1182](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:1182) if self.config.pretraining_tp > 1:

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:795, in Tracer.trace.<locals>.module_call_wrapper(mod, *args, **kwargs)
    [788](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788)     return _orig_module_call(mod, *args, **kwargs)
    [790](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:790) _autowrap_check(
    [791](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:791)     patcher,
    [792](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:792)     getattr(getattr(mod, "forward", mod), "__globals__", {}),
    [793](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:793)     self._autowrap_function_ids,
    [794](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:794) )
--> [795](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:795) return self.call_module(mod, forward, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1170, in HFTracer.call_module(self, m, forward, args, kwargs)
   [1168](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1168)     return forward(*args, **kwargs)
   [1169](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1169) self.orig_forward = forward
-> [1170](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1170) return super().call_module(m, forward, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:479, in Tracer.call_module(self, m, forward, args, kwargs)
    [477](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:477) self.module_stack[_scope.module_path] = _scope.module_type
    [478](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:478) if not self.is_leaf_module(m, module_qualified_name):
--> [479](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:479)     ret_val = forward(*args, **kwargs)
    [480](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:480) else:
    [481](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:481)     ret_val = self.create_proxy("call_module", module_qualified_name, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788, in Tracer.trace.<locals>.module_call_wrapper.<locals>.forward(*args, **kwargs)
    [787](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:787) def forward(*args, **kwargs):
--> [788](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788)     return _orig_module_call(mod, *args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1516](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1516)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1517](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1517) else:
-> [1518](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518)     return self._call_impl(*args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   [1522](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1522) # If we don't have any hooks, we want to skip the rest of the logic in
   [1523](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1523) # this function, and just call forward.
   [1524](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1524) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1525](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1525)         or _global_backward_pre_hooks or _global_backward_hooks
   [1526](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1526)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1527](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1527)     return forward_call(*args, **kwargs)
   [1529](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1529) try:
   [1530](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1530)     result = None

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:969, in LlamaModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
    [958](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:958)     layer_outputs = self._gradient_checkpointing_func(
    [959](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:959)         decoder_layer.__call__,
    [960](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:960)         hidden_states,
   (...)
    [966](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:966)         cache_position,
    [967](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:967)     )
    [968](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:968) else:
--> [969](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:969)     layer_outputs = decoder_layer(
    [970](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:970)         hidden_states,
    [971](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:971)         attention_mask=causal_mask,
    [972](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:972)         position_ids=position_ids,
    [973](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:973)         past_key_value=past_key_values,
    [974](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:974)         output_attentions=output_attentions,
    [975](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:975)         use_cache=use_cache,
    [976](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:976)         cache_position=cache_position,
    [977](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:977)     )
    [979](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:979) hidden_states = layer_outputs[0]
    [981](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:981) if use_cache:

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:795, in Tracer.trace.<locals>.module_call_wrapper(mod, *args, **kwargs)
    [788](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788)     return _orig_module_call(mod, *args, **kwargs)
    [790](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:790) _autowrap_check(
    [791](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:791)     patcher,
    [792](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:792)     getattr(getattr(mod, "forward", mod), "__globals__", {}),
    [793](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:793)     self._autowrap_function_ids,
    [794](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:794) )
--> [795](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:795) return self.call_module(mod, forward, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1170, in HFTracer.call_module(self, m, forward, args, kwargs)
   [1168](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1168)     return forward(*args, **kwargs)
   [1169](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1169) self.orig_forward = forward
-> [1170](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1170) return super().call_module(m, forward, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:479, in Tracer.call_module(self, m, forward, args, kwargs)
    [477](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:477) self.module_stack[_scope.module_path] = _scope.module_type
    [478](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:478) if not self.is_leaf_module(m, module_qualified_name):
--> [479](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:479)     ret_val = forward(*args, **kwargs)
    [480](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:480) else:
    [481](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:481)     ret_val = self.create_proxy("call_module", module_qualified_name, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788, in Tracer.trace.<locals>.module_call_wrapper.<locals>.forward(*args, **kwargs)
    [787](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:787) def forward(*args, **kwargs):
--> [788](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788)     return _orig_module_call(mod, *args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1516](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1516)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1517](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1517) else:
-> [1518](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518)     return self._call_impl(*args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   [1522](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1522) # If we don't have any hooks, we want to skip the rest of the logic in
   [1523](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1523) # this function, and just call forward.
   [1524](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1524) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1525](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1525)         or _global_backward_pre_hooks or _global_backward_hooks
   [1526](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1526)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1527](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1527)     return forward_call(*args, **kwargs)
   [1529](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1529) try:
   [1530](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1530)     result = None

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:714, in LlamaDecoderLayer.forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position)
    [711](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:711) hidden_states = self.input_layernorm(hidden_states)
    [713](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:713) # Self Attention
--> [714](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:714) hidden_states, self_attn_weights, present_key_value = self.self_attn(
    [715](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:715)     hidden_states=hidden_states,
    [716](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:716)     attention_mask=attention_mask,
    [717](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:717)     position_ids=position_ids,
    [718](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:718)     past_key_value=past_key_value,
    [719](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:719)     output_attentions=output_attentions,
    [720](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:720)     use_cache=use_cache,
    [721](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:721)     cache_position=cache_position,
    [722](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:722) )
    [723](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:723) hidden_states = residual + hidden_states
    [725](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:725) # Fully Connected

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:795, in Tracer.trace.<locals>.module_call_wrapper(mod, *args, **kwargs)
    [788](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788)     return _orig_module_call(mod, *args, **kwargs)
    [790](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:790) _autowrap_check(
    [791](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:791)     patcher,
    [792](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:792)     getattr(getattr(mod, "forward", mod), "__globals__", {}),
    [793](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:793)     self._autowrap_function_ids,
    [794](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:794) )
--> [795](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:795) return self.call_module(mod, forward, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1170, in HFTracer.call_module(self, m, forward, args, kwargs)
   [1168](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1168)     return forward(*args, **kwargs)
   [1169](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1169) self.orig_forward = forward
-> [1170](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:1170) return super().call_module(m, forward, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:479, in Tracer.call_module(self, m, forward, args, kwargs)
    [477](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:477) self.module_stack[_scope.module_path] = _scope.module_type
    [478](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:478) if not self.is_leaf_module(m, module_qualified_name):
--> [479](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:479)     ret_val = forward(*args, **kwargs)
    [480](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:480) else:
    [481](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:481)     ret_val = self.create_proxy("call_module", module_qualified_name, args, kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788, in Tracer.trace.<locals>.module_call_wrapper.<locals>.forward(*args, **kwargs)
    [787](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:787) def forward(*args, **kwargs):
--> [788](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/_symbolic_trace.py:788)     return _orig_module_call(mod, *args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1516](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1516)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1517](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1517) else:
-> [1518](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518)     return self._call_impl(*args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   [1522](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1522) # If we don't have any hooks, we want to skip the rest of the logic in
   [1523](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1523) # this function, and just call forward.
   [1524](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1524) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1525](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1525)         or _global_backward_pre_hooks or _global_backward_hooks
   [1526](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1526)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1527](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1527)     return forward_call(*args, **kwargs)
   [1529](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1529) try:
   [1530](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1530)     result = None

File ~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:641, in LlamaSdpaAttention.forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position)
    [636](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:636)     causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
    [638](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:638) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
    [639](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:639) # Reference: https://github.com/pytorch/pytorch/issues/112577.
--> [641](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:641) if query_states.device.type == "cuda" and causal_mask is not None:
    [642](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:642)     query_states = query_states.contiguous()
    [643](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/models/llama/modeling_llama.py:643)     key_states = key_states.contiguous()

File ~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:668, in HFProxy.__bool__(self)
    [666](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:666) if hasattr(self, "_metadata") and self._metadata is not None:
    [667](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:667)     return self._metadata
--> [668](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/transformers/utils/fx.py:668) return super().__bool__()

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:437, in Proxy.__bool__(self)
    [434](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:434)             self.tracer.create_proxy('call_function', assert_fn, (self,), {})
    [435](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:435)             return True
--> [437](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:437) return self.tracer.to_bool(self)

File ~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:300, in TracerBase.to_bool(self, obj)
    [293](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:293) @compatibility(is_backward_compatible=True)
    [294](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:294) def to_bool(self, obj: 'Proxy') -> bool:
    [295](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:295)     """Called when a proxy object is being converted to a boolean, such as
    [296](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:296)     when used in control flow.  Normally we don't know what to do because
    [297](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:297)     we don't know the value of the proxy, but a custom tracer can attach more
    [298](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:298)     information to the graph node using create_node and can choose to return a value.
    [299](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:299)     """
--> [300](https://file+.vscode-resource.vscode-cdn.net/Users/chuhongjie/python/LLM__inference/~/Library/Python/3.9/lib/python/site-packages/torch/fx/proxy.py:300)     raise TraceError('symbolically traced variables cannot be used as inputs to control flow')

TraceError: symbolically traced variables cannot be used as inputs to control flow



Comments:
I guess same as #31200

yes ,the llama and mistral both cause this error,i submit two issue,but no reply ,so do you konw how to fix it? @zucchini-nlp

No, you tagged the right people in issue description. I just wanted to link related issue to each other 

I had the same problem

Hi @Hongjie1Chu @songh11, this should be fixed on the main branch with https://github.com/huggingface/transformers/pull/31574.