Update app.py

app.py

@@ -68,12 +68,12 @@ dict_map = {
     }
 
 tokenizer_vi2en = AutoTokenizer.from_pretrained("vinai/vinai-translate-vi2en-v2", src_lang="vi_VN")
-model_vi2en = AutoModelForSeq2SeqLM.from_pretrained("vinai/vinai-translate-vi2en-v2")
+model_vi2en = AutoModelForSeq2SeqLM.from_pretrained("vinai/vinai-translate-vi2en-v2",device_map="auto")
 
 def translate_vi2en(vi_text: str) -> str:
     for i, j in dict_map.items():
         vi_text = vi_text.replace(i, j)
-    input_ids = tokenizer_vi2en(vi_text, return_tensors="pt").input_ids
+    input_ids = tokenizer_vi2en(vi_text, return_tensors="pt")to("cuda").input_ids
     output_ids = model_vi2en.generate(
         input_ids,
         decoder_start_token_id=tokenizer_vi2en.lang_code_to_id["en_XX"],