fix visualization (#2453)

keras-team · May 31, 2024 · 5cfe7c3 · 5cfe7c3
1 parent 96c3de8
commit 5cfe7c3
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/keras_cv/src/visualization/plot_image_gallery.py b/keras_cv/src/visualization/plot_image_gallery.py
@@ -175,9 +175,6 @@ def plot_image_gallery(
             current_axis.margins(x=0, y=0)
             current_axis.axis("off")
 
-    if path is None and show is None:
-        return fig
-
     if path is not None:
         plt.savefig(
             fname=path,
@@ -190,3 +187,5 @@ def plot_image_gallery(
     elif show:
         plt.show()
         plt.close()
+    else:
+        return fig