fix

huggingface · Jul 30, 2024 · 76cea01 · 76cea01
1 parent c551f03
commit 76cea01
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum_benchmark/backends/llama_cpp/backend.py b/optimum_benchmark/backends/llama_cpp/backend.py
@@ -58,7 +58,7 @@ def prepare_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
         raise ValueError(f"Task {self.config.task} not supported by {self.NAME}")
 
     def forward(self, inputs: Dict[str, Any], kwargs: Dict[str, Any]) -> Any:
-        self.pretrained_model.embed(**inputs["input"])
+        self.pretrained_model.embed(**inputs)
 
     def prefill(self, inputs: Dict[str, Any], kwargs: Dict[str, Any]) -> list[int]:
         next(self.pretrained_model.generate(**inputs))