Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 from peft import AutoPeftModelForCausalLM
 from transformers import AutoTokenizer, GenerationConfig
+import spaces 
 
 # Load model and tokenizer
 model_name = "nafisneehal/Llama-3.2-3B-bnb-4bit-finetuned-TrialBrain-BaselineFeatures-it"
@@ -67,6 +68,7 @@ PRIOR CONCURRENT THERAPY:
 """
 
 # Function to generate response
+@spaces.GPU
 def generate_response(system_instruction, user_input):
     # Prepare the input and move it to the correct device
     inputs = tokenizer([f"### Instruction:\n{system_instruction}\n### Input:\n{user_input}\n### Response:\n"], return_tensors="pt").to(device)