Llama 2 + FSDP Auto Wrap Issue
Author: jmif
Pull request: no

### System Info

- `transformers` version: 4.34.0.dev0
    -  Based off https://github.com/huggingface/transformers/commit/300d6a4a62aac89b3f439110561d5a2268ffad9e
    - One additional patch pulled in to fix an (I think) unrelated issue https://github.com/jmif/transformers/commit/2fe3989ebbb6729e560c6b438b4e1c7ef38412b4
    -  Installing from https://github.com/jmif/transformers/commit/2fe3989ebbb6729e560c6b438b4e1c7ef38412b4 will give you code I'm running
- Platform: Linux-6.2.0-1012-gcp-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: FSDP
        - mixed_precision: fp16
        - use_cpu: False
        - debug: False
        - num_processes: 2
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - fsdp_config: {'fsdp_auto_wrap_policy': 'TRANSFORMER_BASED_WRAP', 'fsdp_backward_prefetch_policy': 'NO_PREFETCH', 'fsdp_forward_prefetch': False, 'fsdp_offload_params': False, 'fsdp_sharding_strategy': 1, 'fsdp_state_dict_type': 'FULL_STATE_DICT', 'fsdp_sync_module_states': True, 'fsdp_use_orig_params': False}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.0.1+cu118 (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?: yes
- Using distributed or parallel set-up in script?: distributed


### Who can help?

@pacman100 @muellerz

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```
import logging

from datasets import load_dataset
from transformers import (
    set_seed,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
    AutoTokenizer,
    AutoModelForSequenceClassification,
)


logger = logging.getLogger()


def train():
    set_seed(42)

    imdb = load_dataset("imdb")

    tokenizer = AutoTokenizer.from_pretrained(
        "meta-llama/Llama-2-7b-hf", max_length=4096
    )

    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForSequenceClassification.from_pretrained(
        "meta-llama/Llama-2-7b-hf",
        num_labels=2,
        max_position_embeddings=1024,
        use_safetensors=True,
    )

    def tokenize(examples):
        return tokenizer(examples["text"], truncation=True)

    train_dataset = imdb["train"].map(tokenize, batched=True)

    auto_find_batch_size = True
    batch_size = 8
    gradient_accumulation_steps = 1

    training_args = TrainingArguments(
        run_name="test",
        output_dir="/tmp/test",
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        auto_find_batch_size=auto_find_batch_size,
        eval_delay=0,
        evaluation_strategy="epoch",
        learning_rate=3e-5,
        weight_decay=3e-6,
        adam_beta1=0.9,
        adam_beta2=0.999,
        adam_epsilon=1e-8,
        num_train_epochs=10,
        lr_scheduler_type="linear",
        warmup_steps=10,
        log_level="info",
        save_strategy="epoch",
        seed=43,
        fp16=True,
        dataloader_drop_last=True,
    )

    collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="longest")

    trainer = Trainer(
        model,
        training_args,
        tokenizer=tokenizer,
        data_collator=collator,
        train_dataset=train_dataset,
    )

    trainer.train()


if __name__ == "__main__":
    train()

```

When I run this script with `auto_find_batch_size = True` I get the following error during model load:

```
Traceback (most recent call last):
  File "llm_training/coarse_relevancy/accelerate/repro.py", line 82, in <module>
    train()
  File "llm_training/coarse_relevancy/accelerate/repro.py", line 78, in train
    trainer.train()
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/trainer.py", line 1556, in train
    return inner_training_loop(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/accelerate/utils/memory.py", line 136, in decorator
    return function(batch_size, *args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/trainer.py", line 1675, in _inner_training_loop
    self.model = self.accelerator.prepare(self.model)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/accelerate/accelerator.py", line 1270, in prepare
    result = tuple(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/accelerate/accelerator.py", line 1271, in <genexpr>
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/accelerate/accelerator.py", line 1083, in _prepare_one
    return self.prepare_model(obj, device_placement=device_placement)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/accelerate/accelerator.py", line 1429, in prepare_model
    model = FSDP(model, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 391, in __init__
    _auto_wrap(auto_wrap_kwargs, fsdp_kwargs, FullyShardedDataParallel)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 55, in _auto_wrap
    raise ValueError(
ValueError: Expected model.layers.0 to NOT be FullyShardedDataParallel if using an `auto_wrap_policy`
```

However, when I set  `auto_find_batch_size = False` I do not. 

### Expected behavior

I do not receive the error.  Note that I'm not yet testing on a machine large enough to fully train the model, so with `auto_find_batch_size = False` I actually get a CUDA OOM.  If I'm reading the stack trace right this indicates I'm progressing past the FSDP exception above.  Including the full CUDA OOM stack trace when setting `auto_find_batch_size = False` in case it helps debug.

```
Traceback (most recent call last):
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 708, in forward
  File "llm-training/llm_training/coarse_relevancy/accelerate/repro.py", line 82, in <module>
    train()
    layer_outputs = decoder_layer(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
  File "llm-training/llm_training/coarse_relevancy/accelerate/repro.py", line 78, in train
    trainer.train()
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/trainer.py", line 1556, in train
    return inner_training_loop(
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 748, in forward
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/trainer.py", line 1838, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/trainer.py", line 2693, in training_step
    loss = self.compute_loss(model, inputs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/trainer.py", line 2718, in compute_loss
    outputs = model(**inputs)
    output = self._fsdp_wrapped_module(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 748, in forward
    output = self._fsdp_wrapped_module(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 959, in forward
    transformer_outputs = self.model(
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 437, in forward
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
    hidden_states = self.mlp(hidden_states)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 708, in forward
    layer_outputs = decoder_layer(
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 748, in forward
    output = self._fsdp_wrapped_module(*args, **kwargs)
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 220, in forward
      File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 424, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 333, in forward
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
    return forward_call(*args, **kwargs)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 114, in forward
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 186, in apply_rotary_pos_emb
    q_embed = (q * cos) + (rotate_half(q) * sin)
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 177, in rotate_half
    return torch.cat((-x2, x1), dim=-1)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 40.00 MiB (GPU 0; 15.77 GiB total capacity; 15.15 GiB already allocated; 1.88 MiB free; 15.25 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
    return F.linear(input, self.weight, self.bias)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 76.00 MiB (GPU 1; 15.77 GiB total capacity; 15.06 GiB already allocated; 37.88 MiB free; 15.21 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

Comments:
Hello, can you try https://github.com/huggingface/accelerate/pull/1919 and let us know if that fixes the issue?

Installed via ` poetry add 'git+https://github.com/pacman100/accelerate.git#fix_fsdp_torch_compile_issue'
` and I get

```
  File ".virtualenvs/llm-training/lib/python3.10/site-packages/transformers/trainer.py", line 1672, in _inner_training_loop
        self.model = self.accelerator.prepare(self.model)self.model = self.accelerator.prepare(self.model)
  File "virtualenvs/llm-training/lib/python3.10/site-packages/accelerate/accelerator.py", line 1221, in prepare
    is_type_fsdp = (type(obj) == FSDP) or isinstance(unwrapped_model, FSDP)
UnboundLocalError: local variable 'unwrapped_model' referenced before assignment
    is_type_fsdp = (type(obj) == FSDP) or isinstance(unwrapped_model, FSDP)
```

Hello, able to reproduce the issue. I believe that when the inner training loop errors out during the backward pass, the model which is wrapped in FSDP unit results in a corrupted state post which the following warning is given before error:
```
/home/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/nn/modules/module.py:2459: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at aten/src/ATen/core/TensorBody.h:489.)
  if p.grad is not None:
```

output logs for `roberta-large` model wherein i print the model before the `accelerator.prepare`. We can see that when it calls the inner loop again, the model is in corrupted state with the internal modules being wrapped in FSDP unit but the whole model not being wrapped in FSDP unit:
```
FullyShardedDataParallel(
  (_fsdp_wrapped_module): RobertaForSequenceClassification(
    (roberta): RobertaModel(
      (embeddings): FullyShardedDataParallel(
        (_fsdp_wrapped_module): RobertaEmbeddings(
          (word_embeddings): Embedding(50265, 1024, padding_idx=1)
          (position_embeddings): Embedding(514, 1024, padding_idx=1)
          (token_type_embeddings): Embedding(1, 1024)
          (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
      (encoder): RobertaEncoder(
        (layer): ModuleList(
          (0-23): 24 x RobertaLayer(
            (attention): RobertaAttention(
              (self): FullyShardedDataParallel(
                (_fsdp_wrapped_module): RobertaSelfAttention(
                  (query): Linear(in_features=1024, out_features=1024, bias=True)
                  (key): Linear(in_features=1024, out_features=1024, bias=True)
                  (value): Linear(in_features=1024, out_features=1024, bias=True)
                  (dropout): Dropout(p=0.1, inplace=False)
                )
              )
              (output): RobertaSelfOutput(
                (dense): Linear(in_features=1024, out_features=1024, bias=True)
                (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
            )
            (intermediate): RobertaIntermediate(
              (dense): Linear(in_features=1024, out_features=4096, bias=True)
              (intermediate_act_fn): GELUActivation()
            )
            (output): RobertaOutput(
              (dense): Linear(in_features=4096, out_features=1024, bias=True)
              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
              (dropout): Dropout(p=0.1, inplace=False)
            )
          )
        )
      )
    )
    (classifier): RobertaClassificationHead(
      (dense): Linear(in_features=1024, out_features=1024, bias=True)
      (dropout): Dropout(p=0.1, inplace=False)
      (out_proj): Linear(in_features=1024, out_features=2, bias=True)
    )
  )
)
***** Running training *****
  Num examples = 25,000
  Num Epochs = 10
  Instantaneous batch size per device = 512
  Training with DataParallel so batch size has been adjusted to: 256
  Total train batch size (w. parallel, distributed & accumulation) = 512
  Gradient Accumulation steps = 1
  Total optimization steps = 480
  Number of trainable parameters = 293,150,210
/home/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/nn/modules/module.py:2459: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at aten/src/ATen/core/TensorBody.h:489.)
  if p.grad is not None:
  0%|                                                                                                  | 0/240 [00:01<?, ?it/s]
/home/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/nn/modules/module.py:2459: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at aten/src/ATen/core/TensorBody.h:489.)
  if p.grad is not None:
in inner training loop
The following columns in the training set don't have a corresponding argument in `RobertaForSequenceClassification.forward` and have been ignored: text. If text are not expected by `RobertaForSequenceClassification.forward`,  you can safely ignore this message.
RobertaForSequenceClassification(
  (roberta): RobertaModel(
    (embeddings): FullyShardedDataParallel(
      (_fsdp_wrapped_module): RobertaEmbeddings(
        (word_embeddings): Embedding(50265, 1024, padding_idx=1)
        (position_embeddings): Embedding(514, 1024, padding_idx=1)
        (token_type_embeddings): Embedding(1, 1024)
        (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (dropout): Dropout(p=0.1, inplace=False)
      )
    )
    (encoder): RobertaEncoder(
      (layer): ModuleList(
        (0-23): 24 x RobertaLayer(
          (attention): RobertaAttention(
            (self): FullyShardedDataParallel(
              (_fsdp_wrapped_module): RobertaSelfAttention(
                (query): Linear(in_features=1024, out_features=1024, bias=True)
                (key): Linear(in_features=1024, out_features=1024, bias=True)
                (value): Linear(in_features=1024, out_features=1024, bias=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
            )
            (output): RobertaSelfOutput(
              (dense): Linear(in_features=1024, out_features=1024, bias=True)
              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
              (dropout): Dropout(p=0.1, inplace=False)
            )
          )
          (intermediate): RobertaIntermediate(
            (dense): Linear(in_features=1024, out_features=4096, bias=True)
            (intermediate_act_fn): GELUActivation()
          )
          (output): RobertaOutput(
            (dense): Linear(in_features=4096, out_features=1024, bias=True)
            (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
        )
      )
    )
  )
  (classifier): RobertaClassificationHead(
    (dense): Linear(in_features=1024, out_features=1024, bias=True)
    (dropout): Dropout(p=0.1, inplace=False)
    (out_proj): Linear(in_features=1024, out_features=2, bias=True)
  )
)
Traceback (most recent call last):
  File "/home/sourab/temp/issue_26019.py", line 82, in <module>
    train()
  File "/home/sourab/temp/issue_26019.py", line 78, in train
    trainer.train()
  File "/home/sourab/transformers/src/transformers/trainer.py", line 1557, in train
    return inner_training_loop(
  File "/home/sourab/accelerate/src/accelerate/utils/memory.py", line 136, in decorator
    return function(batch_size, *args, **kwargs)
  File "/home/sourab/transformers/src/transformers/trainer.py", line 1680, in _inner_training_loop
    self.model = self.accelerator.prepare(self.model)
  File "/home/sourab/accelerate/src/accelerate/accelerator.py", line 1273, in prepare
    result = tuple(
  File "/home/sourab/accelerate/src/accelerate/accelerator.py", line 1274, in <genexpr>
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
  File "/home/sourab/accelerate/src/accelerate/accelerator.py", line 1084, in _prepare_one
    return self.prepare_model(obj, device_placement=device_placement)
  File "/home/sourab/accelerate/src/accelerate/accelerator.py", line 1446, in prepare_model
    model = FSDP(model, **kwargs)
  File "/home/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 463, in __init__
    _auto_wrap(
  File "/home/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 45, in _auto_wrap
    _check_nested_wrapping(root_module)
  File "/home/sourab/miniconda3/envs/hf/lib/python3.10/site-packages/torch/distributed/fsdp/_wrap_utils.py", line 107, in _check_nested_wrapping
    raise ValueError(
ValueError: FSDP auto wrapping requires modules to not already have FSDP applied but found roberta.embeddings in
RobertaForSequenceClassification(
```

I don't see how this can be fixed as it is not touching the transformers/accelerate code when it gets corrupted.

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Is it getting corrupted in torch land?  Happy to open an issue elsewhere but I'm not familiar enough with the internals to understand the source :)

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

used AutoModelForSequenceClassification and 
```
bnb_config = BitsAndBytesConfig(
    load_in_4bit = True,
    bnb_4bit_qunat_type = "nf4",
    bnb_4bit_compute_dtype = torch.float16,
)

```
for loading the model and finetuned this model by using LoRA and saved as "tuned_model"
while loading the model:
```
from transformers import pipeline
pipe = pipeline('text-classification',
                tuned_model,
                device_map="auto")
```
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[<ipython-input-19-fd1b8ca97698>](https://localhost:8080/#) in <cell line: 2>()
      1 from transformers import pipeline
----> 2 pipe = pipeline('text-classification',tuned_model, device_map="auto")

5 frames
[/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py](https://localhost:8080/#) in <dictcomp>(.0)
   3809                 p: {"safetensors_file": f, "weight_name": p, "dtype": str_dtype}
   3810                 for p, f in weight_map.items()
-> 3811                 if param_device_map[p] == "disk"
   3812             }
   3813 

KeyError: 'lm_head.weight'
```
Can any one suggest me how to load this tuned_model ?

Hi @Zuhashaik, this looks like it might be a different problem. Could you open a new issue please?

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.