diff --git a/wyoming_glados/__main__.py b/wyoming_glados/__main__.py index f6e25d7..6d25d81 100644 --- a/wyoming_glados/__main__.py +++ b/wyoming_glados/__main__.py @@ -31,7 +31,7 @@ async def main() -> None: args = parser.parse_args() device_map = { - "rocm": "hip", + "rocm": "cuda", } device = device_map.get(args.device, args.device) diff --git a/wyoming_glados/handler.py b/wyoming_glados/handler.py index e55366a..9ddc745 100644 --- a/wyoming_glados/handler.py +++ b/wyoming_glados/handler.py @@ -48,8 +48,9 @@ def _load_bert_for_language(language: Languages, device: str) -> None: if bert is not None: bert = bert.float() bert.eval() + bert.to(device) bert_models.__loaded_models[language] = bert - _LOGGER.info("BERT model for %s cast to float32", language.name) + _LOGGER.info("BERT model for %s cast to float32 and moved to %s", language.name, device) def _find_model_files(model_dir: Path):