Skip to content

Commit

Permalink
Update sft_trainer.py (#1241)
Browse files Browse the repository at this point in the history
  • Loading branch information
younesbelkada authored Jan 17, 2024
1 parent bcccdeb commit 4fb3d0c
Showing 1 changed file with 1 addition and 5 deletions.
6 changes: 1 addition & 5 deletions trl/trainer/sft_trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -424,11 +424,7 @@ def tokenize(element):
else:
self._dataset_sanity_checked = True

return {
"input_ids": outputs["input_ids"],
"labels": outputs["input_ids"],
"attention_mask": outputs["attention_mask"],
}
return {"input_ids": outputs["input_ids"], "attention_mask": outputs["attention_mask"]}

signature_columns = ["input_ids", "labels", "attention_mask"]

Expand Down

0 comments on commit 4fb3d0c

Please sign in to comment.