diff --git a/src/training/plot_metrics.py b/src/training/plot_metrics.py
index c8db1aa..10fdbc4 100644
--- a/src/training/plot_metrics.py
+++ b/src/training/plot_metrics.py
@@ -162,7 +162,9 @@ def plot_combined(train_df: pd.DataFrame, val_df: pd.DataFrame, out_dir: Path) -
 
     # 1. Train + Val loss.
     ax = axes[0]
-    sns.lineplot(data=train_df, x="epoch", y="total", ax=ax, marker="o", linewidth=2, color="crimson", label="train")
+    train_loss_col = "train_loss" if "train_loss" in train_df.columns else "total"
+    if train_loss_col in train_df.columns:
+        sns.lineplot(data=train_df, x="epoch", y=train_loss_col, ax=ax, marker="o", linewidth=2, color="crimson", label="train")
     if "loss" in val_df.columns:
         sns.lineplot(data=val_df, x="epoch", y="loss", ax=ax, marker="s", linewidth=2, color="royalblue", label="val")
     ax.set_title("Loss (train vs val)")