Optimization for compute_influence part

logix-project · Mar 31, 2024 · ade2632 · ade2632
1 parent 127efd7
commit ade2632
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/analog/analysis/influence_function.py b/analog/analysis/influence_function.py
@@ -149,12 +149,8 @@ def compute_influence(
         )
 
         # Assign total_influence values to the corresponding locations
-        src_indices = [
-            self.influence_scores.index.get_loc(src_id) for src_id in src_ids
-        ]
-        tgt_indices = [
-            self.influence_scores.columns.get_loc(tgt_id) for tgt_id in tgt_ids
-        ]
+        src_indices = self.influence_scores.index.get_indexer(src_ids)
+        tgt_indices = self.influence_scores.columns.get_indexer(tgt_ids)
 
         self.influence_scores.iloc[src_indices, tgt_indices] = total_influence.numpy()