fix tgi

huggingface · Dec 10, 2024 · 1bddc66 · 1bddc66
1 parent 46bfe63
commit 1bddc66
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum_benchmark/backends/py_txi/config.py b/optimum_benchmark/backends/py_txi/config.py
@@ -50,7 +50,7 @@ class PyTXIConfig(BackendConfig):
     quantize: Optional[str] = None
     num_shard: Optional[int] = None
     speculate: Optional[int] = None
-    cuda_graphs: Optional[bool] = None
+    cuda_graphs: Optional[int] = None
     disable_custom_kernels: Optional[bool] = None
     trust_remote_code: Optional[bool] = None