Bug Fix: Check for explicit cli device (fast) (pytorch#1374)

Jack-Khuu · web-flow · commit 46977645de6e · 2024-11-13T18:04:34.000-08:00
diff --git a/torchchat/cli/cli.py b/torchchat/cli/cli.py
@@ -533,7 +533,7 @@ def arg_init(args):
         # Localized import to minimize expensive imports
         from torchchat.utils.build_utils import get_device_str
 
-        if args.device is None:
+        if args.device is None or args.device == "fast":
             args.device = get_device_str(
                 args.quantize.get("executor", {}).get("accelerator", default_device)
             )

Original file line number	Diff line number	Diff line change
`@@ -533,7 +533,7 @@ def arg_init(args):`
`533`	`533`	`# Localized import to minimize expensive imports`
`534`	`534`	`from torchchat.utils.build_utils import get_device_str`
`535`	`535`
`536`		`- if args.device is None:`
	`536`	`+ if args.device is None or args.device == "fast":`
`537`	`537`	`args.device = get_device_str(`
`538`	`538`	`args.quantize.get("executor", {}).get("accelerator", default_device)`
`539`	`539`	`)`