diff --git a/plot_single_transmission_timeline.py b/plot_single_transmission_timeline.py
index 904a024..69db340 100755
--- a/plot_single_transmission_timeline.py
+++ b/plot_single_transmission_timeline.py
@@ -122,7 +122,7 @@ if __name__ == "__main__":
 
         cmap = matplotlib.cm.get_cmap("Set3")
         unique_cells = transmission_df["cell_color"].unique()
-        color_list = cmap * (round(len(unique_cells) / len(cmap)) + 1)
+        color_list = cmap.colors * (round(len(unique_cells) / len(cmap.colors)) + 1)
 
         for c in transmission_df["cell_color"].unique():
             bounds = transmission_df[["index", "cell_color"]].groupby("cell_color").agg(["min", "max"]).loc[c]
diff --git a/plot_transmission_timeline.py b/plot_transmission_timeline.py
index d6d54fa..ea7c45e 100755
--- a/plot_transmission_timeline.py
+++ b/plot_transmission_timeline.py
@@ -206,7 +206,7 @@ if __name__ == "__main__":
 
     cmap = matplotlib.cm.get_cmap("Set3")
     unique_cells = transmission_df["cell_color"].unique()
-    color_list = cmap * (round(len(unique_cells) / len(cmap)) + 1)
+    color_list = cmap.colors * (round(len(unique_cells) / len(cmap.colors)) + 1)
 
     for c in transmission_df["cell_color"].unique():
         bounds = transmission_df[["index", "cell_color"]].groupby("cell_color").agg(["min", "max"]).loc[c]