diff --git a/aimsim/__init__.py b/aimsim/__init__.py
index e16ecd1..877e143 100644
--- a/aimsim/__init__.py
+++ b/aimsim/__init__.py
@@ -7,4 +7,4 @@
 except ImportError:
     pass  # aimsim_core does not include this
 
-__version__ = "2.2.0"
+__version__ = "2.2.1"
diff --git a/aimsim/tasks/cluster_data.py b/aimsim/tasks/cluster_data.py
index 034f735..40aafb4 100644
--- a/aimsim/tasks/cluster_data.py
+++ b/aimsim/tasks/cluster_data.py
@@ -3,6 +3,10 @@
 from os.path import dirname
 
 import matplotlib.pyplot as plt
+try:
+    from matplotlib.cm import get_cmap
+except:  # noqa: E722
+    from matplotlib.pyplot import get_cmap
 from matplotlib.colors import rgb2hex
 import yaml
 
@@ -54,7 +58,7 @@ def _extract_configs(self):
         self.plot_settings = dict()
         self.plot_settings["cluster_plot"] = {
             "cluster_colors": [
-                rgb2hex(plt.cm.get_cmap("tab20", self.n_clusters)(cluster_id))
+                rgb2hex(get_cmap("tab20", self.n_clusters)(cluster_id))
                 for cluster_id in range(self.n_clusters)],
             "response": "Response",
         }