From 7c77d1c5b72a0d6513d1ee721882abeafef80814 Mon Sep 17 00:00:00 2001 From: "l.gabrysiak" Date: Tue, 25 Feb 2025 17:00:56 +0100 Subject: [PATCH] mod --- hft.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hft.py b/hft.py index 6190912..9348bae 100644 --- a/hft.py +++ b/hft.py @@ -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)