update

src/attribution/attntrace.py

@@ -36,6 +36,7 @@ class AttnTraceAttribution(Attribution):
                 #print(last_group,group, last_group_label,group_label)
                 importances[feature_index[0]]+=(last_group_loss-group_loss)
         return importances
+    @spaces.GPU
     def attribute(self, question: str, contexts: list, answer: str,explained_answer: str, customized_template: str = None):
         start_time = time.time()
         if self.llm.model!=None:

src/models/Llama.py

@@ -34,6 +34,7 @@ class Llama(Model):
             self.model = model
         return self.model
 
+    @spaces.GPU
     def query(self, msg, max_tokens=128000):
         model = self._load_model_if_needed().to("cuda")
         messages = self.messages