.

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM
 import torch
 import logging
 
@@ -13,7 +13,7 @@ logger = logging.getLogger(__name__)
 
 # Predefined list of models to compare (can be expanded)
 model_options = {
-    "Foundation-Sec-8B": pipeline("text-generation", model="fdtn-ai/Foundation-Sec-8B"),
+    "Foundation-Sec-8B": pipeline("text-generation", model="fdtn-ai/Foundation-Sec-8B", torch_dtype=torch.bfloat16),
 }
 
 @spaces.GPU