🚧 [WIP/TEST] Model : CUDAでの処理実装 part 1

2024-10-18 16:16:42 +09:00
parent 147b5cc396
commit 5ba595cefe
2 changed files with 3 additions and 3 deletions
--- a/src-python/models/translation/translation_translator.py
+++ b/src-python/models/translation/translation_translator.py
@@ -37,7 +37,7 @@ class Translator():
        tokenizer_path = os_path.join(path, "weights", "ctranslate2", directory_name, "tokenizer")
        self.ctranslate2_translator = ctranslate2.Translator(
            weight_path,
-            device="cpu",
+            device="cuda",
            device_index=0,
            compute_type="int8",
            inter_threads=1,