Add n_gpu_layers

app.py

@@ -71,7 +71,9 @@ def generate_text_from_prompt(user_prompt,
                               top_p=0.8,
                               echo=True,
                               stop=["<s>", "[/INST]", "Q"]):
-    mistral_model = Llama(model_path=model_file_path, n_ctx=4096)
+    mistral_model = Llama(model_path=model_file_path,
+                          n_ctx=4096,
+                          n_gpu_layers=24)
 
     # Define the parameters
     model_output = mistral_model(