Add reward to image caption

This commit is contained in:
Kevin Black 2023-06-28 10:42:47 -07:00
parent fe9ed8a25f
commit ad28862b48

View File

@ -263,8 +263,8 @@ def main(_):
accelerator.log( accelerator.log(
{ {
"images": [ "images": [
wandb.Image(os.path.join(tmpdir, f"{i}.jpg"), caption=prompt) wandb.Image(os.path.join(tmpdir, f"{i}.jpg"), caption=f"{prompt:.25} | {reward:.2f}")
for i, prompt in enumerate(prompts) for i, (prompt, reward) in enumerate(zip(prompts, rewards))
], ],
}, },
step=global_step, step=global_step,