From e46a89e2bc05800d2e4432e79ac9a2265d70b16e Mon Sep 17 00:00:00 2001
From: TheDudeFromCI <thedudefromci@gmail.com>
Date: Wed, 31 Aug 2022 02:19:12 -0700
Subject: [PATCH] Update step before training checkpoint

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

diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index 6cc82ef..884f81f 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -846,6 +846,7 @@ class Trainer(object):
 
                 accelerator.wait_for_everyone()
 
+                self.step += 1
                 if accelerator.is_main_process:
                     self.ema.to(device)
                     self.ema.update()
@@ -862,7 +863,6 @@ class Trainer(object):
                         utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = int(math.sqrt(self.num_samples)))
                         self.save(milestone)
 
-                self.step += 1
                 pbar.update(1)
 
         accelerator.print('training complete')