From ad28862b48bf2265cad80dd517bd8ff29441ef34 Mon Sep 17 00:00:00 2001
From: Kevin Black <kevin3.black@gmail.com>
Date: Wed, 28 Jun 2023 10:42:47 -0700
Subject: [PATCH] Add reward to image caption

---
 scripts/train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scripts/train.py b/scripts/train.py
index 58395d9..6e7214c 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -263,8 +263,8 @@ def main(_):
             accelerator.log(
                 {
                     "images": [
-                        wandb.Image(os.path.join(tmpdir, f"{i}.jpg"), caption=prompt)
-                        for i, prompt in enumerate(prompts)
+                        wandb.Image(os.path.join(tmpdir, f"{i}.jpg"), caption=f"{prompt:.25} | {reward:.2f}")
+                        for i, (prompt, reward) in enumerate(zip(prompts, rewards))
                     ],
                 },
                 step=global_step,