Merge pull request #45 from int-brain-lab/patch_wfl

emerg WFL patch
int-brain-lab · Sep 18, 2024 · 2ff5528 · 2ff5528
2 parents a28b8cc + e76cd23
commit 2ff5528
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/ibldsp/waveform_extraction.py b/src/ibldsp/waveform_extraction.py
@@ -584,7 +584,7 @@ def load_waveforms(self, labels=None, indices=None, return_info=True, flatten=Fa
             indices = np.arange(self.max_wf)
 
         labels = np.array(labels)
-        label_idx = np.array([np.where(labels == label)[0][0] for label in labels])
+        label_idx = np.array([np.where(self.labels == label)[0][0] for label in labels])
         indices = np.array(indices)
 
         num_labels = labels.shape[0]
@@ -650,7 +650,7 @@ def random_waveforms(
             assert num_random_labels is None, "labels and num_random_labels cannot both be set"
 
         labels = np.array(labels)
-        label_idx = np.array([np.where(labels == label)[0][0] for label in labels])
+        label_idx = np.array([np.where(self.labels == label)[0][0] for label in labels])
 
         num_labels = labels.shape[0]