From 800dabd1fdc8d9672518b12f30bbeae34c200e84 Mon Sep 17 00:00:00 2001
From: ADKoishi <2574994844@qq.com>
Date: Thu, 18 Jan 2024 05:29:47 +0800
Subject: [PATCH 1/2] Fixed fatal bug where model failed to save

---
 scripts/train_with_rm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/train_with_rm.py b/scripts/train_with_rm.py
index c821b4a..2ddc29e 100644
--- a/scripts/train_with_rm.py
+++ b/scripts/train_with_rm.py
@@ -572,7 +572,7 @@ def compare(a, b):
                 # make sure we did an optimization step at the end of the inner epoch
                 assert accelerator.sync_gradients
 
-            if epoch!=0 and (epoch+1) % config.save_freq == 0 and accelerator.is_main_process:
+            if epoch!=0 and (epoch+1) % config.save_freq == 0:
                 accelerator.save_state()
 
 

From bbb8245d9dc3c42ecac5ef0ffddd7963ab3621d8 Mon Sep 17 00:00:00 2001
From: ADKoishi <2574994844@qq.com>
Date: Thu, 18 Jan 2024 20:27:25 +0800
Subject: [PATCH 2/2] Fixed saving failure with bf16 AMP.

---
 scripts/train_with_rm.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/scripts/train_with_rm.py b/scripts/train_with_rm.py
index 2ddc29e..3c91c20 100644
--- a/scripts/train_with_rm.py
+++ b/scripts/train_with_rm.py
@@ -387,7 +387,8 @@ def compare(a, b):
         # this is a hack to force wandb to log the images as JPEGs instead of PNGs
         with tempfile.TemporaryDirectory() as tmpdir:
             for i, image in enumerate(images):
-                pil = Image.fromarray((image[0].cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8))
+                image_np = image[0].to('cpu').to(torch.float32).numpy()
+                pil = Image.fromarray((image_np.transpose(1, 2, 0) * 255).astype(np.uint8))
                 pil = pil.resize((256, 256))
                 pil.save(os.path.join(tmpdir, f"{i}.jpg"))
             accelerator.log(