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

This commit is contained in:
misyaguziya
2024-10-18 16:16:42 +09:00
parent 147b5cc396
commit 5ba595cefe
2 changed files with 3 additions and 3 deletions

View File

@@ -51,7 +51,7 @@ def checkWhisperWeight(root, weight_type):
try: try:
WhisperModel( WhisperModel(
path, path,
device="cpu", device="cuda",
device_index=0, device_index=0,
compute_type="int8", compute_type="int8",
cpu_threads=4, cpu_threads=4,
@@ -79,7 +79,7 @@ def getWhisperModel(root, weight_type):
path = os_path.join(root, "weights", "whisper", weight_type) path = os_path.join(root, "weights", "whisper", weight_type)
return WhisperModel( return WhisperModel(
path, path,
device="cpu", device="cuda",
device_index=0, device_index=0,
compute_type="int8", compute_type="int8",
cpu_threads=4, cpu_threads=4,

View File

@@ -37,7 +37,7 @@ class Translator():
tokenizer_path = os_path.join(path, "weights", "ctranslate2", directory_name, "tokenizer") tokenizer_path = os_path.join(path, "weights", "ctranslate2", directory_name, "tokenizer")
self.ctranslate2_translator = ctranslate2.Translator( self.ctranslate2_translator = ctranslate2.Translator(
weight_path, weight_path,
device="cpu", device="cuda",
device_index=0, device_index=0,
compute_type="int8", compute_type="int8",
inter_threads=1, inter_threads=1,