for lucidrains#36

andrewmoise · Jun 11, 2022 · ecc6f30 · ecc6f30
1 parent f900f40
commit ecc6f30
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
@@ -7,6 +7,7 @@
 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 @@ def __init__(
         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
@@ -3,7 +3,7 @@
 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',