From 6f0d4387337d66b61df1041523f7307076b17b49 Mon Sep 17 00:00:00 2001
From: Sarthak Jariwala <jariwala@uw.edu>
Date: Thu, 7 Dec 2023 07:56:13 -0800
Subject: [PATCH] use ax.get_figure to fetch figure

---
 src/seaborn_image/_core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/seaborn_image/_core.py b/src/seaborn_image/_core.py
index 2029f2c3..35329672 100644
--- a/src/seaborn_image/_core.py
+++ b/src/seaborn_image/_core.py
@@ -69,7 +69,7 @@ def _setup_figure(self):
         if self.ax is None:
             f, ax = plt.subplots()
         else:
-            f = plt.gcf()
+            f = self.ax.get_figure() or plt.gcf()
             ax = self.ax
 
         return f, ax