diff --git a/Inst_eddies/Analysis/plot_velocity.py b/Inst_eddies/Analysis/plot_velocity.py
index 7b343d53fbf992300a05d222983e55f65da69b7a..88a9f08b282207389650b31f6a957e6d462d55d6 100755
--- a/Inst_eddies/Analysis/plot_velocity.py
+++ b/Inst_eddies/Analysis/plot_velocity.py
@@ -89,12 +89,7 @@ def plot_velocity(args):
 
     ax.gridlines(draw_labels=True)
     ax.coastlines()
-
-    if args.save:
-        fig.savefig(f"plot_velocity.{args.save}")
-        print(f'Created "plot_velocity.{args.save}".')
-    else:
-        plt.show()
+    return fig
 
 
 if __name__ == "__main__":
@@ -129,4 +124,10 @@ if __name__ == "__main__":
     )
     parser.add_argument("input_file", help="NetCDF file containing velocity")
     args = parser.parse_args()
-    plot_velocity(args)
+    fig = plot_velocity(args)
+
+    if args.save:
+        fig.savefig(f"plot_velocity.{args.save}")
+        print(f'Created "plot_velocity.{args.save}".')
+    else:
+        plt.show()