diff --git a/loopstructuralvisualisation/_3d_viewer.py b/loopstructuralvisualisation/_3d_viewer.py
index 7714c51..ff00a67 100644
--- a/loopstructuralvisualisation/_3d_viewer.py
+++ b/loopstructuralvisualisation/_3d_viewer.py
@@ -383,7 +383,7 @@ def _build_stratigraphic_cmap(self, model):
                     data.append((v["id"], v["colour"]))
                     colours.append(v["colour"])
                     boundaries.append(v["id"])  # print(u,v)
-            cmap = colors.ListedColormap(colours).colors
+            cmap = colors.ListedColormap(colours)
         except ImportError:
             logger.warning("Cannot use predefined colours as I can't import matplotlib")
             cmap = "tab20"
@@ -433,7 +433,7 @@ def plot_model_surfaces(
             strati_surfaces = []
             surfaces = model.get_stratigraphic_surfaces()
             if cmap is None:
-                cmap = self._build_stratigraphic_cmap(model)
+                cmap = self._build_stratigraphic_cmap(model).colors
             for s in surfaces:
                 strati_surfaces.append(s.vtk())
             if name is None: