Update app.py

app.py

@@ -10,6 +10,7 @@ from transformers import (
 from peft import PeftModel, PeftConfig
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 peft_model_id = "Namkoy/whisper_peft_vi_nam"
 language = "vietnamese"
@@ -17,7 +18,7 @@ task = "transcribe"
 peft_config = PeftConfig.from_pretrained(peft_model_id)
 model = WhisperForConditionalGeneration.from_pretrained(
     peft_config.base_model_name_or_path, load_in_8bit=True, device_map="auto"
-)
+).to(device)
 
 model = PeftModel.from_pretrained(model, peft_model_id)
 tokenizer = WhisperTokenizer.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)