Update app.py

app.py

@@ -15,7 +15,7 @@ def load_models():
 
     for call_name, real_name in model_name_dict.items():
         print('\tLoading model: %s' % call_name)
-        model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(real_name).cuda()
         tokenizer = AutoTokenizer.from_pretrained(real_name)
         model_dict[call_name+'_model'] = model
         model_dict[call_name+'_tokenizer'] = tokenizer
@@ -35,7 +35,7 @@ def translation(source, target, text):
     model = model_dict[model_name + '_model']
     tokenizer = model_dict[model_name + '_tokenizer']
 
-    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
+    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target, device=0)
     output = translator(text, max_length=400)
 
     end_time = time.time()