diff --git a/src/spikeinterface/core/template_tools.py b/src/spikeinterface/core/template_tools.py
index 934b18ed49..3c8663df70 100644
--- 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: