update

src/attribution/attntrace.py

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