Update main.py

main.py

@@ -54,11 +54,14 @@ def convert_to_16_bit_wav(data):
         raise ValueError("Audio data cannot be converted to " "16-bit int format.")
     return data
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
 # Load the model with INT8 quantization
 model = AutoModelForCausalLM.from_pretrained(
     model_path,
     cache_dir=".",
-    load_in_8bit=True,  # Enable loading in INT8
+    load_in_8bit=True if 'cuda' in device else False,  # Enable loading in INT8
     device_map="auto"  # Automatically map model to available devices
 )
 
@@ -67,7 +70,6 @@ config_path = "audiotokenizer/speechtokenizer_hubert_avg_config.json"
 ckpt_path = "audiotokenizer/SpeechTokenizer.pt"
 quantizer = SpeechTokenizer.load_from_checkpoint(config_path, ckpt_path)
 quantizer.eval()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Перемещение всех слоев квантизатора на устройство и их заморозка
 def freeze_entire_model(model):