update

src/attribution/attntrace.py

@@ -37,13 +37,17 @@ 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 load_model(self):
+        self.model = self.llm._load_model_if_needed().to("cuda")
+    
     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:
             self.model = self.llm.model
         else:
-            self.model = self.llm._load_model_if_needed().to("cuda")
+            self.load_model()
         self.layers = range(len(self.model.model.layers))
         model = self.model
         tokenizer = self.tokenizer