Refactor code

app.py

@@ -12,7 +12,7 @@ model = torch.load('saved_model.pth', map_location= torch.device(device), weight
 #     else:
 #         idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
     
-#     yield model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature)[0].tolist() 
+#     yield model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature) 
     
 
 
@@ -35,6 +35,6 @@ with gr.Blocks() as demo:
         idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
     generate_btn = gr.Button(value="Generate")
     outputs = [gr.Textbox(label= "Generated text: ")]
-    generate_btn.click(fn = model.generate(idx, max_new_tokens=num_of_tokens,temperature=tmp)[0].tolist() , inputs= inputs, outputs= outputs)
+    generate_btn.click(fn = model.generate, inputs= inputs, outputs= outputs)
     
 demo.launch()