placed input_ids in cuda and just added the vector for input_ids on generate() for model

app.py

@@ -34,9 +34,10 @@ def llama_generation(input_text):
     """
 
     input_ids = llama_tokenizer.encode(input_text,
-                                       return_tensors='pt')
+                                       return_tensors='pt').to('cuda')
     
-    output_ids = llama_model.generate(**input_ids)
+    # llama generation looks for the numeric vectors not the tensors so there is no need for **input_ids rather just input_ids
+    output_ids = llama_model.generate(input_ids=input_ids) 
 
     # Decode
     output_text = llama_tokenizer.decode(output_ids,