Merge pull request SpikeInterface#3495 from OlivierPeron/Loading-temp…

…lates Template loading average or median
JoeZiminski · Oct 23, 2024 · 4a1a45a · 4a1a45a
2 parents bfe9fb6 + 72357a6
commit 4a1a45a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/spikeinterface/core/template_tools.py b/src/spikeinterface/core/template_tools.py
@@ -31,7 +31,12 @@ def get_dense_templates_array(one_object: Templates | SortingAnalyzer, return_sc
             )
         ext = one_object.get_extension("templates")
         if ext is not None:
-            templates_array = ext.data["average"]
+            if "average" in ext.data:
+                templates_array = ext.data.get("average")
+            elif "median" in ext.data:
+                templates_array = ext.data.get("median")
+            else:
+                raise ValueError("Average or median templates have not been computed.")
         else:
             raise ValueError("SortingAnalyzer need extension 'templates' to be computed to retrieve templates")
     else: