Skip to content

Commit 615954a

Browse files
committed
Default to CUDA if available (#431)
Unless user specifies device, default it to CUDA if it's available on the platform. As for all models, other than TinyLlama inference on GPU will be faster than on CPU
1 parent bbd4463 commit 615954a

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

build/builder.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,16 @@ class BuilderArgs:
3232
gguf_kwargs: Optional[Dict[str, Any]] = None
3333
dso_path: Optional[Union[Path, str]] = None
3434
pte_path: Optional[Union[Path, str]] = None
35-
device: str = "cpu"
35+
device: Optional[str] = None
3636
precision: torch.dtype = torch.float32
3737
setup_caches: bool = False
3838
use_tp: bool = False
3939
is_chat_model: bool = False
4040

4141
def __post_init__(self):
42+
if self.device is None:
43+
self.device = "cuda" if torch.cuda.is_available() else "cpu"
44+
4245
if not (
4346
(self.checkpoint_path and self.checkpoint_path.is_file())
4447
or (self.checkpoint_dir and self.checkpoint_dir.is_dir())

0 commit comments

Comments
 (0)