[generate] generate not working gradient_checkpointing=True
Author: alcompa
Pull request: no
Labels: name: bug

name: Generation

### System Info

```
transformers==4.37.2 
peft==0.13.1 
accelerate==0.21.0 
deepspeed==0.12.6
torch==2.1.2+cu118
```

### Who can help?

@zucchini-nlp @ArthurZucker  

### 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

## Minimal reproducible example on colab
```bash
!pip install -q uv
!uv pip install -q --system transformers==4.37.2 peft==0.13.1 accelerate==0.21.0 deepspeed==0.12.6
!pip install -q torch==2.1.2 --index-url https://download.pytorch.org/whl/cu118
```

```python
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer,
    AutoProcessor
)
from torch.utils.data import Dataset
import torch
import os


class DummyDataset(Dataset):
    def __init__(self):
        self.samples = [{"text": "Hello world"} for _ in range(10)]

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx]

model = AutoModelForCausalLM.from_pretrained(
    "HuggingFaceM4/tiny-random-LlamaForCausalLM",
    torch_dtype=torch.float16
)
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceM4/tiny-random-LlamaForCausalLM")

def collate_fn(batch):
    texts = [item["text"] for item in batch]
    return tokenizer(texts, return_tensors="pt", padding=True, truncation=True)


# A dummy trainer in which we need to generate during the training loop
class MyTrainer(Seq2SeqTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        with torch.no_grad():
            # simulating generation with inputs-embeds (like in Llava)
            # https://github.com/haotian-liu/LLaVA/blob/c121f0432da27facab705978f83c4ada465e46fd/llava/model/language_model/llava_llama.py#L137-L141
            generated = model.generate(
                inputs_embeds=torch.randn((4, 100, 16), dtype=torch.float16, device=self.args.device),
                attention_mask=torch.where(torch.arange(100).unsqueeze(0).repeat_interleave(4, dim=0) < 50, torch.zeros((1,)), torch.ones((1,))).to(dtype=torch.bool, device=self.args.device),
                do_sample=True,
                temperature=0.6,
                output_scores=False,
                num_return_sequences=2,
                max_new_tokens=30,
            )

        print("Generated:", tokenizer.decode(generated[0]))
        return torch.tensor(0.5, requires_grad=True, device=self.args.device)

training_args = Seq2SeqTrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=2,
    learning_rate=1e-5,
    num_train_epochs=1,
    logging_steps=1,
    gradient_checkpointing=True, # True # <-- Generation works with grad. ckpt. disabled, doesn't work with grad. ckpt. enabled
    remove_unused_columns=False,
    fp16=True,
    optim="adamw_torch",
    report_to="none"
)

trainer = MyTrainer(
    model=model,
    args=training_args,
    train_dataset=DummyDataset(),
    data_collator=collate_fn,
    tokenizer=tokenizer,
)

trainer.train()
```

## Bug description
`generate()` in `transformers==4.37.2` does not work properly when `gradient_checkpointing` is enabled in `Seq2SeqTrainer`.

If `gradient_checkpointing=False`, `generate()` works fine and returns the generated sequence: 
```
Generated: <unk> radical助 reflected pipapplejnapom workaroundчик moment coresкра información Lisaвали résocoAnim radical ore Gene after convergence allem FIFAtembre davon
```

With `gradient_checkpointing=True` generation raises the error `RuntimeError: The size of tensor a (101) must match the size of tensor b (100) at non-singleton dimension 3`

## Full traceback
```
warnings.warn(
/usr/local/lib/python3.11/dist-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  return self.fget.__get__(instance, owner)()
`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...
/usr/local/lib/python3.11/dist-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.
  warnings.warn(
/usr/local/lib/python3.11/dist-packages/torch/utils/checkpoint.py:61: UserWarning: None of the inputs have requires_grad=True. Gradients will be None
  warnings.warn(
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-4-ab8901164bef>](https://localhost:8080/#) in <cell line: 0>()
     72 )
     73 
---> 74 trainer.train()

17 frames
[/usr/local/lib/python3.11/dist-packages/transformers/trainer.py](https://localhost:8080/#) in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1537                 hf_hub_utils.enable_progress_bars()
   1538         else:
-> 1539             return inner_training_loop(
   1540                 args=args,
   1541                 resume_from_checkpoint=resume_from_checkpoint,

[/usr/local/lib/python3.11/dist-packages/transformers/trainer.py](https://localhost:8080/#) in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1867 
   1868                 with self.accelerator.accumulate(model):
-> 1869                     tr_loss_step = self.training_step(model, inputs)
   1870 
   1871                 if (

[/usr/local/lib/python3.11/dist-packages/transformers/trainer.py](https://localhost:8080/#) in training_step(self, model, inputs)
   2770 
   2771         with self.compute_loss_context_manager():
-> 2772             loss = self.compute_loss(model, inputs)
   2773 
   2774         if self.args.n_gpu > 1:

[<ipython-input-4-ab8901164bef>](https://localhost:8080/#) in compute_loss(self, model, inputs, return_outputs, **kwargs)
     38             # simulating generation with inputs-embeds (like in Llava)
     39             # https://github.com/haotian-liu/LLaVA/blob/c121f0432da27facab705978f83c4ada465e46fd/llava/model/language_model/llava_llama.py#L137-L141
---> 40             generated = model.generate(
     41                 inputs_embeds=torch.randn((4, 100, 16), dtype=torch.float16, device=self.args.device),
     42                 attention_mask=torch.where(torch.arange(100).unsqueeze(0).repeat_interleave(4, dim=0) < 50, torch.zeros((1,)), torch.ones((1,))).to(dtype=torch.bool, device=self.args.device),

[/usr/local/lib/python3.11/dist-packages/torch/utils/_contextlib.py](https://localhost:8080/#) in decorate_context(*args, **kwargs)
    113     def decorate_context(*args, **kwargs):
    114         with ctx_factory():
--> 115             return func(*args, **kwargs)
    116 
    117     return decorate_context

[/usr/local/lib/python3.11/dist-packages/transformers/generation/utils.py](https://localhost:8080/#) in generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   1523 
   1524             # 13. run sample
-> 1525             return self.sample(
   1526                 input_ids,
   1527                 logits_processor=prepared_logits_processor,

[/usr/local/lib/python3.11/dist-packages/transformers/generation/utils.py](https://localhost:8080/#) in sample(self, input_ids, logits_processor, stopping_criteria, logits_warper, max_length, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, streamer, **model_kwargs)
   2620 
   2621             # forward pass to get next token
-> 2622             outputs = self(
   2623                 **model_inputs,
   2624                 return_dict=True,

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.11/dist-packages/accelerate/utils/operations.py](https://localhost:8080/#) in forward(*args, **kwargs)
    579 
    580     def forward(*args, **kwargs):
--> 581         return model_forward(*args, **kwargs)
    582 
    583     # To act like a decorator so that it can be popped when doing `extract_model_from_parallel`

[/usr/local/lib/python3.11/dist-packages/accelerate/utils/operations.py](https://localhost:8080/#) in __call__(self, *args, **kwargs)
    567 
    568     def __call__(self, *args, **kwargs):
--> 569         return convert_to_fp32(self.model_forward(*args, **kwargs))
    570 
    571     def __getstate__(self):

[/usr/local/lib/python3.11/dist-packages/torch/amp/autocast_mode.py](https://localhost:8080/#) in decorate_autocast(*args, **kwargs)
     14     def decorate_autocast(*args, **kwargs):
     15         with autocast_instance:
---> 16             return func(*args, **kwargs)
     17 
     18     decorate_autocast.__script_unsupported = "@autocast() decorator is not supported in script mode"  # type: ignore[attr-defined]

[/usr/local/lib/python3.11/dist-packages/transformers/models/llama/modeling_llama.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1181 
   1182         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1183         outputs = self.model(
   1184             input_ids=input_ids,
   1185             attention_mask=attention_mask,

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.11/dist-packages/transformers/models/llama/modeling_llama.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
   1033             # output_attentions=True can not be supported when using SDPA, and we fall back on
   1034             # the manual implementation that requires a 4D causal mask in all cases.
-> 1035             attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
   1036                 attention_mask,
   1037                 (batch_size, seq_length),

[/usr/local/lib/python3.11/dist-packages/transformers/modeling_attn_mask_utils.py](https://localhost:8080/#) in _prepare_4d_causal_attention_mask_for_sdpa(attention_mask, input_shape, inputs_embeds, past_key_values_length, sliding_window)
    396         )
    397     else:
--> 398         expanded_4d_mask = attn_mask_converter.to_4d(
    399             attention_mask,
    400             input_shape[-1],

[/usr/local/lib/python3.11/dist-packages/transformers/modeling_attn_mask_utils.py](https://localhost:8080/#) in to_4d(self, attention_mask_2d, query_length, dtype, key_value_length)
    135 
    136         if causal_4d_mask is not None:
--> 137             expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
    138 
    139         # expanded_attn_mask + causal_4d_mask can cause some overflow

RuntimeError: The size of tensor a (101) must match the size of tensor b (100) at non-singleton dimension 3
```

### Expected behavior

I think this error is caused by the fact that after the first model forward, `_update_model_kwargs_for_generation` extends the attention mask for the new generated token
https://github.com/huggingface/transformers/blob/345b9b1a6a308a1fa6559251eb33ead2211240ac/src/transformers/generation/utils.py#L623

but `LlamaForCausalLM.prepare_inputs_for_generation` always relies on `inputs_embeds` to generate the next token when `gradient_checkpointing=True --> use_cache=False --> past_key_values=None`
https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py#L1270-L1273

That does not happen when gradient_checkpointing=False because past_key_values is not None after the first model forward so input_ids are chosen as model_inputs for successive steps.

I can't tell if this affects also recent version of transformers because I am not familiar with the new Cache API.

Is it enough to find an alternative way to check whether we are in the first generation step or in a successive one (instead of relying on past_key_values) to fix the issue or there are other functions that relies on this check which I am not aware of?

Thanks in advance for the support :)

Comments:
@alcompa hey! Can you try to update to the latest version and see if the error persists? We've changed a lot since v4.37 regarding cache and generation, so might have been resolved already



I tried to run the same snippet (gradient_checkpointing=True) with `transformers==4.48.2` and `accelerate==1.3.0` and I get a slightly different error: `RuntimeError: Expected size for first two dimensions of batch2 tensor to be: [32, 101] but got: [32, 1].`
but it's still related to sdpa attention computation.
Also in this case generate works properly when gradient checkpointing is disabled

EDIT: same results with transformers main branch

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[4], [line 74](vscode-notebook-cell:?execution_count=4&line=74)
     [53](vscode-notebook-cell:?execution_count=4&line=53) training_args = Seq2SeqTrainingArguments(
     [54](vscode-notebook-cell:?execution_count=4&line=54)     output_dir="./output",
     [55](vscode-notebook-cell:?execution_count=4&line=55)     per_device_train_batch_size=2,
   (...)
     [63](vscode-notebook-cell:?execution_count=4&line=63)     report_to="none"
     [64](vscode-notebook-cell:?execution_count=4&line=64) )
     [66](vscode-notebook-cell:?execution_count=4&line=66) trainer = MyTrainer(
     [67](vscode-notebook-cell:?execution_count=4&line=67)     model=model,
     [68](vscode-notebook-cell:?execution_count=4&line=68)     args=training_args,
   (...)
     [71](vscode-notebook-cell:?execution_count=4&line=71)     tokenizer=tokenizer,
     [72](vscode-notebook-cell:?execution_count=4&line=72) )
---> [74](vscode-notebook-cell:?execution_count=4&line=74) trainer.train()

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2171, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   [2169](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2169)         hf_hub_utils.enable_progress_bars()
   [2170](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2170) else:
-> [2171](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2171)     return inner_training_loop(
   [2172](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2172)         args=args,
   [2173](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2173)         resume_from_checkpoint=resume_from_checkpoint,
   [2174](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2174)         trial=trial,
   [2175](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2175)         ignore_keys_for_eval=ignore_keys_for_eval,
   [2176](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2176)     )

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2531, in Trainer._inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   [2524](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2524) context = (
   [2525](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2525)     functools.partial(self.accelerator.no_sync, model=model)
   [2526](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2526)     if i != len(batch_samples) - 1
   [2527](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2527)     and self.accelerator.distributed_type != DistributedType.DEEPSPEED
   [2528](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2528)     else contextlib.nullcontext
   [2529](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2529) )
   [2530](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2530) with context():
-> [2531](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2531)     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
   [2533](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2533) if (
   [2534](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2534)     args.logging_nan_inf_filter
   [2535](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2535)     and not is_torch_xla_available()
   [2536](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2536)     and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
   [2537](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2537) ):
   [2538](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2538)     # if loss is nan or inf simply add the average of previous logged losses
   [2539](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:2539)     tr_loss = tr_loss + tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3675, in Trainer.training_step(self, model, inputs, num_items_in_batch)
   [3672](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3672)     return loss_mb.reduce_mean().detach().to(self.args.device)
   [3674](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3674) with self.compute_loss_context_manager():
-> [3675](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3675)     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
   [3677](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3677) del inputs
   [3678](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3678) if (
   [3679](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3679)     self.args.torch_empty_cache_steps is not None
   [3680](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3680)     and self.state.global_step % self.args.torch_empty_cache_steps == 0
   [3681](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/trainer.py:3681) ):

Cell In[4], [line 40](vscode-notebook-cell:?execution_count=4&line=40)
     [36](vscode-notebook-cell:?execution_count=4&line=36) def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
     [37](vscode-notebook-cell:?execution_count=4&line=37)     with torch.no_grad():
     [38](vscode-notebook-cell:?execution_count=4&line=38)         # simulating generation with inputs-embeds (like in Llava)
     [39](vscode-notebook-cell:?execution_count=4&line=39)         # https://github.com/haotian-liu/LLaVA/blob/c121f0432da27facab705978f83c4ada465e46fd/llava/model/language_model/llava_llama.py#L137-L141
---> [40](vscode-notebook-cell:?execution_count=4&line=40)         generated = model.generate(
     [41](vscode-notebook-cell:?execution_count=4&line=41)             inputs_embeds=torch.randn((4, 100, 16), dtype=torch.float16, device=self.args.device),
     [42](vscode-notebook-cell:?execution_count=4&line=42)             attention_mask=torch.where(torch.arange(100).unsqueeze(0).repeat_interleave(4, dim=0) < 50, torch.zeros((1,)), torch.ones((1,))).to(dtype=torch.bool, device=self.args.device),
     [43](vscode-notebook-cell:?execution_count=4&line=43)             do_sample=True,
     [44](vscode-notebook-cell:?execution_count=4&line=44)             temperature=0.6,
     [45](vscode-notebook-cell:?execution_count=4&line=45)             output_scores=False,
     [46](vscode-notebook-cell:?execution_count=4&line=46)             num_return_sequences=2,
     [47](vscode-notebook-cell:?execution_count=4&line=47)             max_new_tokens=30,
     [48](vscode-notebook-cell:?execution_count=4&line=48)         )
     [50](vscode-notebook-cell:?execution_count=4&line=50)     print("Generated:", tokenizer.decode(generated[0]))
     [51](vscode-notebook-cell:?execution_count=4&line=51)     return torch.tensor(0.5, requires_grad=True, device=self.args.device)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    [112](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py:112) @functools.wraps(func)
    [113](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py:113) def decorate_context(*args, **kwargs):
    [114](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py:114)     with ctx_factory():
--> [115](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_contextlib.py:115)         return func(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2255, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   [2247](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2247)     input_ids, model_kwargs = self._expand_inputs_for_generation(
   [2248](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2248)         input_ids=input_ids,
   [2249](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2249)         expand_size=generation_config.num_return_sequences,
   [2250](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2250)         is_encoder_decoder=self.config.is_encoder_decoder,
   [2251](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2251)         **model_kwargs,
   [2252](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2252)     )
   [2254](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2254)     # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
-> [2255](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2255)     result = self._sample(
   [2256](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2256)         input_ids,
   [2257](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2257)         logits_processor=prepared_logits_processor,
   [2258](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2258)         stopping_criteria=prepared_stopping_criteria,
   [2259](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2259)         generation_config=generation_config,
   [2260](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2260)         synced_gpus=synced_gpus,
   [2261](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2261)         streamer=streamer,
   [2262](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2262)         **model_kwargs,
   [2263](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2263)     )
   [2265](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2265) elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH):
   [2266](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2266)     # 11. prepare beam search scorer
   [2267](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2267)     beam_scorer = BeamSearchScorer(
   [2268](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2268)         batch_size=batch_size,
   [2269](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2269)         num_beams=generation_config.num_beams,
   (...)
   [2274](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2274)         max_length=generation_config.max_length,
   [2275](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:2275)     )

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3257, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, **model_kwargs)
   [3255](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3255)     is_prefill = False
   [3256](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3256) else:
-> [3257](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3257)     outputs = model_forward(**model_inputs, return_dict=True)
   [3259](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3259) # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
   [3260](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3260) model_kwargs = self._update_model_kwargs_for_generation(
   [3261](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3261)     outputs,
   [3262](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3262)     model_kwargs,
   [3263](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3263)     is_encoder_decoder=self.config.is_encoder_decoder,
   [3264](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/generation/utils.py:3264) )

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1509](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1509)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1510](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1510) else:
-> [1511](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511)     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
   [1515](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1515) # If we don't have any hooks, we want to skip the rest of the logic in
   [1516](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1516) # this function, and just call forward.
   [1517](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1517) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1518](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518)         or _global_backward_pre_hooks or _global_backward_hooks
   [1519](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1519)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1520](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520)     return forward_call(*args, **kwargs)
   [1522](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1522) try:
   [1523](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1523)     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/accelerate/utils/operations.py:819, in convert_outputs_to_fp32.<locals>.forward(*args, **kwargs)
    [818](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/accelerate/utils/operations.py:818) def forward(*args, **kwargs):
--> [819](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/accelerate/utils/operations.py:819)     return model_forward(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/accelerate/utils/operations.py:807, in ConvertOutputsToFp32.__call__(self, *args, **kwargs)
    [806](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/accelerate/utils/operations.py:806) def __call__(self, *args, **kwargs):
--> [807](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/accelerate/utils/operations.py:807)     return convert_to_fp32(self.model_forward(*args, **kwargs))

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/amp/autocast_mode.py:16, in autocast_decorator.<locals>.decorate_autocast(*args, **kwargs)
     [13](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/amp/autocast_mode.py:13) @functools.wraps(func)
     [14](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/amp/autocast_mode.py:14) def decorate_autocast(*args, **kwargs):
     [15](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/amp/autocast_mode.py:15)     with autocast_instance:
---> [16](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/amp/autocast_mode.py:16)         return func(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:831, 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, num_logits_to_keep, **kwargs)
    [828](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:828) return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    [830](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:830) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
--> [831](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:831) outputs = self.model(
    [832](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:832)     input_ids=input_ids,
    [833](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:833)     attention_mask=attention_mask,
    [834](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:834)     position_ids=position_ids,
    [835](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:835)     past_key_values=past_key_values,
    [836](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:836)     inputs_embeds=inputs_embeds,
    [837](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:837)     use_cache=use_cache,
    [838](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:838)     output_attentions=output_attentions,
    [839](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:839)     output_hidden_states=output_hidden_states,
    [840](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:840)     return_dict=return_dict,
    [841](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:841)     cache_position=cache_position,
    [842](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:842)     **kwargs,
    [843](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:843) )
    [845](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:845) hidden_states = outputs[0]
    [846](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:846) # Only compute necessary logits, and do not upcast them to float if we are not computing the loss

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1509](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1509)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1510](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1510) else:
-> [1511](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511)     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
   [1515](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1515) # If we don't have any hooks, we want to skip the rest of the logic in
   [1516](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1516) # this function, and just call forward.
   [1517](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1517) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1518](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518)         or _global_backward_pre_hooks or _global_backward_hooks
   [1519](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1519)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1520](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520)     return forward_call(*args, **kwargs)
   [1522](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1522) try:
   [1523](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1523)     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:577, 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, **flash_attn_kwargs)
    [574](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:574)     all_hidden_states += (hidden_states,)
    [576](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:576) if self.gradient_checkpointing and self.training:
--> [577](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:577)     layer_outputs = self._gradient_checkpointing_func(
    [578](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:578)         decoder_layer.__call__,
    [579](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:579)         hidden_states,
    [580](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:580)         causal_mask,
    [581](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:581)         position_ids,
    [582](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:582)         past_key_values,
    [583](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:583)         output_attentions,
    [584](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:584)         use_cache,
    [585](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:585)         cache_position,
    [586](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:586)         position_embeddings,
    [587](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:587)     )
    [588](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:588) else:
    [589](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:589)     layer_outputs = decoder_layer(
    [590](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:590)         hidden_states,
    [591](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:591)         attention_mask=causal_mask,
   (...)
    [598](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:598)         **flash_attn_kwargs,
    [599](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:599)     )

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_compile.py:24, in _disable_dynamo.<locals>.inner(*args, **kwargs)
     [20](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_compile.py:20) @functools.wraps(fn)
     [21](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_compile.py:21) def inner(*args, **kwargs):
     [22](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_compile.py:22)     import torch._dynamo
---> [24](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_compile.py:24)     return torch._dynamo.disable(fn, recursive)(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:489, in _TorchDynamoContext.__call__.<locals>._fn(*args, **kwargs)
    [487](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:487)     dynamo_config_ctx.__enter__()
    [488](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:488) try:
--> [489](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:489)     return fn(*args, **kwargs)
    [490](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:490) finally:
    [491](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:491)     set_eval_frame(prior)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:17, in wrap_inline.<locals>.inner(*args, **kwargs)
     [15](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:15) @functools.wraps(fn)
     [16](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:16) def inner(*args, **kwargs):
---> [17](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:17)     return fn(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:482, in checkpoint(function, use_reentrant, context_fn, determinism_check, debug, *args, **kwargs)
    [477](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:477)     if context_fn is not noop_context_fn or debug is not False:
    [478](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:478)         raise ValueError(
    [479](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:479)             "Passing `context_fn` or `debug` is only supported when "
    [480](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:480)             "use_reentrant=False."
    [481](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:481)         )
--> [482](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:482)     return CheckpointFunction.apply(function, preserve, *args)
    [483](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:483) else:
    [484](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:484)     gen = _checkpoint_without_reentrant_generator(
    [485](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:485)         function, preserve, context_fn, determinism_check, debug, *args, **kwargs
    [486](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:486)     )

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:553, in Function.apply(cls, *args, **kwargs)
    [550](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:550) if not torch._C._are_functorch_transforms_active():
    [551](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:551)     # See NOTE: [functorch vjp and autograd interaction]
    [552](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:552)     args = _functorch.utils.unwrap_dead_wrappers(args)
--> [553](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:553)     return super().apply(*args, **kwargs)  # type: ignore[misc]
    [555](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:555) if not is_setup_ctx_defined:
    [556](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:556)     raise RuntimeError(
    [557](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:557)         "In order to use an autograd.Function with functorch transforms "
    [558](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:558)         "(vmap, grad, jvp, jacrev, ...), it must override the setup_context "
    [559](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:559)         "staticmethod. For more details, please see "
    [560](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:560)         "https://pytorch.org/docs/master/notes/extending.func.html"
    [561](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/function.py:561)     )

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:261, in CheckpointFunction.forward(ctx, run_function, preserve_rng_state, *args)
    [258](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:258) ctx.save_for_backward(*tensor_inputs)
    [260](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:260) with torch.no_grad():
--> [261](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:261)     outputs = run_function(*args)
    [262](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/checkpoint.py:262) return outputs

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1509](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1509)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1510](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1510) else:
-> [1511](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511)     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
   [1515](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1515) # If we don't have any hooks, we want to skip the rest of the logic in
   [1516](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1516) # this function, and just call forward.
   [1517](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1517) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1518](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518)         or _global_backward_pre_hooks or _global_backward_hooks
   [1519](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1519)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1520](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520)     return forward_call(*args, **kwargs)
   [1522](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1522) try:
   [1523](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1523)     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:332, in LlamaDecoderLayer.forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position, position_embeddings, **kwargs)
    [329](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:329) hidden_states = self.input_layernorm(hidden_states)
    [331](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:331) # Self Attention
--> [332](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:332) hidden_states, self_attn_weights = self.self_attn(
    [333](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:333)     hidden_states=hidden_states,
    [334](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:334)     attention_mask=attention_mask,
    [335](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:335)     position_ids=position_ids,
    [336](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:336)     past_key_value=past_key_value,
    [337](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:337)     output_attentions=output_attentions,
    [338](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:338)     use_cache=use_cache,
    [339](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:339)     cache_position=cache_position,
    [340](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:340)     position_embeddings=position_embeddings,
    [341](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:341)     **kwargs,
    [342](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:342) )
    [343](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:343) hidden_states = residual + hidden_states
    [345](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:345) # Fully Connected

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1509](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1509)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1510](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1510) else:
-> [1511](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1511)     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
   [1515](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1515) # If we don't have any hooks, we want to skip the rest of the logic in
   [1516](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1516) # this function, and just call forward.
   [1517](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1517) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1518](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518)         or _global_backward_pre_hooks or _global_backward_hooks
   [1519](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1519)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1520](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1520)     return forward_call(*args, **kwargs)
   [1522](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1522) try:
   [1523](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1523)     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:288, in LlamaAttention.forward(self, hidden_states, position_embeddings, attention_mask, past_key_value, cache_position, **kwargs)
    [285](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:285)     else:
    [286](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:286)         attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
--> [288](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:288) attn_output, attn_weights = attention_interface(
    [289](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:289)     self,
    [290](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:290)     query_states,
    [291](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:291)     key_states,
    [292](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:292)     value_states,
    [293](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:293)     attention_mask,
    [294](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:294)     dropout=0.0 if not self.training else self.attention_dropout,
    [295](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:295)     scaling=self.scaling,
    [296](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:296)     **kwargs,
    [297](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:297) )
    [299](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:299) attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    [300](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:300) attn_output = self.o_proj(attn_output)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:53, in sdpa_attention_forward(module, query, key, value, attention_mask, dropout, scaling, is_causal, **kwargs)
     [50](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:50) if torch.jit.is_tracing() and isinstance(is_causal, torch.Tensor):
     [51](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:51)     is_causal = is_causal.item()
---> [53](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:53) attn_output = torch.nn.functional.scaled_dot_product_attention(
     [54](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:54)     query,
     [55](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:55)     key,
     [56](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:56)     value,
     [57](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:57)     attn_mask=causal_mask,
     [58](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:58)     dropout_p=dropout,
     [59](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:59)     scale=scaling,
     [60](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:60)     is_causal=is_causal,
     [61](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:61) )
     [62](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:62) attn_output = attn_output.transpose(1, 2).contiguous()
     [64](https://vscode-remote+vscode-002d01jkbbh7w2q0q1s5ck5y2c28xn-002estudio-002elightning-002eai.vscode-resource.vscode-cdn.net/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:64) return attn_output, None

RuntimeError: Expected size for first two dimensions of batch2 tensor to be: [32, 101] but got: [32, 1].
```

Ah I see now, you are trying to generate with `inputs_embeds` and using grad checkpointing. But GC is not compatible with cache while generating from `inputs_embeds` works only with cache

I guess you'd have to choose only one option, and we might need to raise a more informative error for these cases

Thanks, I missed this detail.
Another thing I noticed is that when I generate with `input_ids`, `GenerateDecoderOnlyOutput.sequences` contains prompt+completion, instead when generating when inputs_embeds `GenerateDecoderOnlyOutput.sequences` contains only the completion. Is this expected (because `inputs_embeds` can't be "converted" back to `input_ids`)?

Exactly, expected behavior and you pointed out the reason correctly. AFAIR we had a plan to enable users to remove prompt, with smth like `return_full_text`, but not in very near future

Thanks :)

Closing as resolved then