Repository: huggingface/trl
Issue category: Bug Report

Issue:
Wrong tensor index for roll and truncate in DPOTrainer fn concatenated_forward( ).
### System Info

it is a tensor index error

### Information

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

### Tasks

- [ ] An officially supported task in the `examples` folder
- [ ] My own task or dataset (give details below)

### Reproduction

```python
# Get the first column idx that is all zeros and remove every column after that
empty_cols = torch.sum(attention_mask, dim=0) == 0
first_empty_col = torch.nonzero(empty_cols)[0].item() if empty_cols.any() else attention_mask.size(1) + 1
input_ids = input_ids[:, : first_empty_col - 1]
attention_mask = attention_mask[:, : first_empty_col - 1]
loss_mask = loss_mask[:, : first_empty_col - 1]
```

### Expected behavior

The returns of _torch.nonzero_ is the index (starts from 0) of non-zero elements, so there is no need to add -1 to _first_empty_col_. 
The correct code should be:
```python
empty_cols = torch.sum(attention_mask, dim=0) == 0
first_empty_col = torch.nonzero(empty_cols)[0].item() if empty_cols.any() else attention_mask.size(1)
input_ids = input_ids[:, : first_empty_col]
attention_mask = attention_mask[:, : first_empty_col]
loss_mask = loss_mask[:, : first_empty_col]
```


Hints:
Good catch! Thanks! Do you mind opening a PR to fix that?

Patch:
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
index b563cab2f5..0c9883387a 100644
--- a/trl/trainer/dpo_trainer.py
+++ b/trl/trainer/dpo_trainer.py
@@ -1086,10 +1086,10 @@ def concatenated_forward(self, model: nn.Module, batch: Dict[str, Union[List, to
 
             # Get the first column idx that is all zeros and remove every column after that
             empty_cols = torch.sum(attention_mask, dim=0) == 0
-            first_empty_col = torch.nonzero(empty_cols)[0].item() if empty_cols.any() else attention_mask.size(1) + 1
-            input_ids = input_ids[:, : first_empty_col - 1]
-            attention_mask = attention_mask[:, : first_empty_col - 1]
-            loss_mask = loss_mask[:, : first_empty_col - 1]
+            first_empty_col = torch.nonzero(empty_cols)[0].item() if empty_cols.any() else attention_mask.size(1)
+            input_ids = input_ids[:, : first_empty_col]
+            attention_mask = attention_mask[:, : first_empty_col]
+            loss_mask = loss_mask[:, : first_empty_col]
 
             # Truncate right
             if self.args.max_length is not None:


Test patch: