diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 1c9fbc9..781fe6f 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -7,6 +7,7 @@ from inspect import isfunction from functools import partial from torch.utils import data +from multiprocessing import cpu_count from torch.cuda.amp import autocast, GradScaler from pathlib import Path @@ -598,7 +599,7 @@ class Trainer(object): self.train_num_steps = train_num_steps self.ds = Dataset(folder, image_size, augment_horizontal_flip = augment_horizontal_flip) - self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle=True, pin_memory=True)) + self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count())) self.opt = Adam(diffusion_model.parameters(), lr=train_lr) self.step = 0 diff --git a/setup.py b/setup.py index 955cc9d..3cbb45f 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.18.2', + version = '0.18.3', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',