fixed image classification input generator

huggingface · Jan 13, 2024 · 99df40d · 99df40d
1 parent f93acee
commit 99df40d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum_benchmark/generators/task_generator.py b/optimum_benchmark/generators/task_generator.py
@@ -290,7 +290,7 @@ def labels(self):
 
     def generate(self):
         dummy = {}
-        dummy["pixel_value"] = self.pixel_values()
+        dummy["pixel_values"] = self.pixel_values()
 
         if self.with_labels:
             dummy["labels"] = self.labels()
@@ -318,7 +318,7 @@ def labels(self):
 
     def generate(self):
         dummy = {}
-        dummy["pixel_value"] = self.pixel_values()
+        dummy["pixel_values"] = self.pixel_values()
 
         if self.with_labels:
             dummy["labels"] = self.labels()