diff --git a/contact_map/plot_utils.py b/contact_map/plot_utils.py
index a373c72..3e3ee02 100644
--- a/contact_map/plot_utils.py
+++ b/contact_map/plot_utils.py
@@ -22,6 +22,7 @@
     'afmhot', 'gist_heat', 'copper'
 ]
 
+
 def is_cmap_diverging(cmap):
     if cmap in _DIVERGING:
         return True
@@ -31,6 +32,7 @@ def is_cmap_diverging(cmap):
         warnings.warn("Unknown colormap: Treating as sequential.")
         return False
 
+
 def ranged_colorbar(cmap, norm, cbmin, cbmax, ax=None):
     """Create a colorbar with given endpoints.
 
@@ -60,6 +62,7 @@ def ranged_colorbar(cmap, norm, cbmin, cbmax, ax=None):
 
     if ax is None:
         fig = plt
+        ax = plt.gca()
     else:
         fig = ax.figure