Fix namespace error of dtype args (#139)

huggingface · Jan 15, 2025 · 84c7fb8 · 84c7fb8
1 parent fe5e7c8
commit 84c7fb8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/benchmark_pipelines.py b/scripts/benchmark_pipelines.py
@@ -12,11 +12,11 @@
 
 def get_transformers_pipeline(args: Namespace):
     if "dtype" in args:
-        assert args["dtype"] in {"float16", "bfloat16", "float32"}
+        assert args.dtype in {"float16", "bfloat16", "float32"}
 
     return raw_pipeline(
         model=args.model,
-        torch_dtype=args["dtype"],
+        torch_dtype=args.dtype,
         model_kwargs={
             "device_map": "balanced",
             "max_memory": {0: "20GiB", "cpu": "64GiB"},