This commit is contained in:
l.gabrysiak 2025-02-25 17:00:56 +01:00
parent c0077c7c04
commit 7c77d1c5b7
1 changed files with 1 additions and 1 deletions

2
hft.py
View File

@ -139,7 +139,7 @@ class CustomModel(AutoModelForCausalLM):
return outputs
class CustomTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
labels = inputs.pop("labels")
source_idx = inputs.pop("source_idx", None)
outputs = model(**inputs, labels=labels, source_idx=source_idx)