fix

huggingface · Jan 14, 2024 · 074ce58 · 074ce58
1 parent e9ff880
commit 074ce58
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/optimum_benchmark/backends/onnxruntime/backend.py b/optimum_benchmark/backends/onnxruntime/backend.py
@@ -381,11 +381,11 @@ def prepare_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
             return {"prompt": inputs["prompt"]}
 
         LOGGER.info(f"\t+ Moving inputs tensors to device {self.device}")
-        for key, value in inputs.items():
-            if key not in self.inputs_names:
-                inputs.pop(key)
-            else:
+        for key, value in list(inputs.items()):
+            if key in self.inputs_names:
                 inputs[key] = value.to(self.device)
+            else:
+                inputs.pop(key)
 
         return inputs