Update app.py

app.py

@@ -1,3 +1,6 @@
+import torch
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import streamlit as st
 
@@ -17,8 +20,8 @@ sentence = st.text_input('Input your sentence here:', value='My favorite ice cre
 
 st.info("Max generated sentence: 100 words")
 if (st.button("Generate")):
-    input_ids = tokenizer.encode(sentence, return_tensors='pt')
-    paragraph_generated = model.generate(input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True)
-    text = tokenizer.decode(paragraph_generated[0], skip_special_tokens=True)
+    input_ids = tokenizer.encode(sentence, return_tensors='pt').to(device)
+    paragraph_generated = model.generate(input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True).to(device)
+    text = tokenizer.decode(paragraph_generated[0], skip_special_tokens=True).to(device)
 
     st.write(text)