Bug fix: Enable fast to override quantize json (#1377)

* Bug fix: Enable fast to override quantize json * collapse conditional
pytorch · Nov 20, 2024 · b809b69 · b809b69
1 parent 8dccc5a
commit b809b69
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/torchchat/cli/cli.py b/torchchat/cli/cli.py
@@ -533,16 +533,16 @@ def arg_init(args):
         # Localized import to minimize expensive imports
         from torchchat.utils.build_utils import get_device_str
 
-        if args.device is None or args.device == "fast":
+        if args.device is None:
             args.device = get_device_str(
                 args.quantize.get("executor", {}).get("accelerator", default_device)
             )
         else:
+            args.device = get_device_str(args.device)
             executor_handler = args.quantize.get("executor", None)
-            if executor_handler:
-                if executor_handler["accelerator"] != args.device:
-                    print('overriding json-specified device {executor_handler["accelerator"]} with cli device {args.device}')
-                    executor_handler["accelerator"] = args.device
+            if executor_handler and executor_handler["accelerator"] != args.device:
+                print(f'overriding json-specified device {executor_handler["accelerator"]} with cli device {args.device}')
+                executor_handler["accelerator"] = args.device
 
     if "mps" in args.device:
         if getattr(args, "compile", False) or getattr(args, "compile_prefill", False):