Update app.py

app.py

@@ -7,17 +7,17 @@ urlretrieve("https://github.com/equ1/generative_python_transformer/tree/main/GPT
 
 # inference function
 def inference(inp):
-    tokenizer = AutoTokenizer.from_pretrained("GPT-python")
-    model = AutoModelWithLMHead.from_pretrained("GPT-python")
+    tokenizer = AutoTokenizer.from_pretrained("GPT-python").to("cuda")
+    model = AutoModelWithLMHead.from_pretrained("GPT-python").to("cuda")
 
-    input_ids = tokenizer.encode(inp, return_tensors="pt")
+    input_ids = tokenizer.encode(inp, return_tensors="pt").to("cuda")
     beam_output = model.generate(input_ids, 
                                max_length=512,
                                num_beams=10,
                                temperature=0.7,
                                no_repeat_ngram_size=5,
                                num_return_sequences=1,
-                               )
+                               ).to("cuda")
   
     output = []
     for beam in beam_output: