chore: fix new api

zama-ai · Jan 10, 2025 · 2891a9e · 2891a9e
1 parent d2a401d
commit 2891a9e
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/use_case_examples/lora_finetuning/eval.py b/use_case_examples/lora_finetuning/eval.py
@@ -206,7 +206,11 @@ def metric_fn(model, dataloader):
     0, tokenizer.vocab_size, (PER_DEVICE_TRAIN_BATCH_SIZE, BLOCK_SIZE), dtype=torch.long
 )
 attention_mask = torch.ones((PER_DEVICE_TRAIN_BATCH_SIZE, BLOCK_SIZE), dtype=torch.long)
-inputset = (input_tensor, label_tensor, attention_mask)
+inputset = {
+    "input_ids": input_tensor,
+    "attention_mask": attention_mask,
+    "labels": label_tensor
+}
 
 # Prepare eval loader
 eval_loader = DataLoader(test_dataset, batch_size=4, shuffle=False, collate_fn=data_collator)