From 5698fdc7216ac6ba220e68d4df92d277e73f0184 Mon Sep 17 00:00:00 2001
From: Bear <bearc@polaris-login-02.hsn.cm.polaris.alcf.anl.gov>
Date: Fri, 21 Mar 2025 18:50:08 +0000
Subject: [PATCH] Move training legend outside figure for train visualization

---
 src/spine/vis/train.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/spine/vis/train.py b/src/spine/vis/train.py
index 72a33453..93ef9ad0 100644
--- a/src/spine/vis/train.py
+++ b/src/spine/vis/train.py
@@ -421,7 +421,12 @@ def draw(
                 plt.ylabel(ylabel if len(metric) == 1 else "Metric")
                 plt.gca().set_ylim(limits[metric[0]])
                 legend_title = model_name[model[0]] if len(model) == 1 else None
-                plt.legend(ncol=leg_ncols, title=legend_title)
+                plt.legend(
+                    ncol=leg_ncols,
+                    title=legend_title,
+                    bbox_to_anchor=(1.05, 1),
+                    loc="upper left",
+                )
 
             if figure_name:
                 plt.savefig(f"{figure_name}.png", bbox_inches="tight")