fix accelerator prepare bug for dataloader

andrewmoise · Jul 25, 2022 · 6b504c4 · 6b504c4
1 parent 37334ae
commit 6b504c4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -719,6 +719,7 @@ def __init__(
         self.ds = Dataset(folder, self.image_size, augment_horizontal_flip = augment_horizontal_flip, convert_image_to = convert_image_to)
         dl = DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count())
 
+        dl = self.accelerator.prepare(dl)
         self.dl = cycle(dl)
 
         # optimizer
@@ -739,7 +740,7 @@ def __init__(
 
         # prepare model, dataloader, optimizer with accelerator
 
-        self.model, self.dl, self.opt = self.accelerator.prepare(self.model, self.dl, self.opt)
+        self.model, self.opt = self.accelerator.prepare(self.model, self.opt)
 
     def save(self, milestone):
         if not self.accelerator.is_local_main_process:

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.26.1',
+  version = '0.26.3',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',