Update src/litbench_pipeline.py

src/litbench_pipeline.py

@@ -33,7 +33,7 @@ import torch
 from peft.peft_model import PeftModel
 from datasets import load_dataset
 from huggingface_hub import login
-
+import spaces
 
 os.environ["CUDA_VISIBLE_DEVICES"]="0"
 
@@ -311,6 +311,7 @@ def fetch_arxiv_papers(papers_to_download):
 
 
 # Chat prediction function
+@spaces.GPU
 def predict(message, history, selected_task):
     global model
     # Initialize the conversation string