From 6b504c4ae9bffa6c36cbfb6a23ee5aba11c4e508 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Mon, 25 Jul 2022 08:13:46 -0700 Subject: [PATCH] fix accelerator prepare bug for dataloader --- denoising_diffusion_pytorch/denoising_diffusion_pytorch.py | 3 ++- setup.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index ef751cc..31212b0 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -719,6 +719,7 @@ class Trainer(object): 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 @@ class Trainer(object): # 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 index 8a3b245..ccfd88a 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages 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',