Add reward to image caption

2023-06-28 10:42:47 -07:00
parent fe9ed8a25f
commit ad28862b48
1 changed files with 2 additions and 2 deletions
--- 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,