fix another stray bug

README.md

@@ -1,6 +1,6 @@
 <img src="./denoising-diffusion.png" width="500px"></img>
 
-## Denoising Diffusion Probabilistic Model, in Pytorch (wip)
+## Denoising Diffusion Probabilistic Model, in Pytorch
 
 Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch. It is a new approach to generative modeling that may <a href="https://ajolicoeur.wordpress.com/the-new-contender-to-gans-score-matching-with-langevin-sampling/">have the potential</a> to rival GANs. It uses denoising score matching to estimate the gradient of the data distribution, followed by Langevin sampling to sample from the true distribution. This implementation was transcribed from the official Tensorflow version <a href="https://github.com/hojonathanho/diffusion">here</a>. 
 
@@ -61,9 +61,9 @@ trainer = Trainer(
     'path/to/your/images',
     image_size = 128,
     train_batch_size = 32,
-    train_lr = 3e-4,
+    train_lr = 2e-5,
     train_num_steps = 100000,
-    gradient_accumulate_every = 1
+    gradient_accumulate_every = 2
 )
 
 trainer.train()

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -372,9 +372,9 @@ class Trainer(object):
         *,
         image_size = 128,
         train_batch_size = 32,
-        train_lr = 3e-4,
+        train_lr = 2e-5,
         train_num_steps = 100000,
-        gradient_accumulate_every = 1
+        gradient_accumulate_every = 2
     ):
         super().__init__()
         self.model = diffusion_model
@@ -403,7 +403,7 @@ class Trainer(object):
                 milestone = ind // SAVE_AND_SAMPLE_EVERY
                 all_images = self.model.p_sample_loop((64, 3, self.image_size, self.image_size))
                 utils.save_image(all_images, f'./sample-{milestone}.png', nrow=8)
-                torch.save(model.state_dict(), f'./model-{milestone}.pt')
+                torch.save(self.model.state_dict(), f'./model-{milestone}.pt')
 
             ind += 1
 

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.1.0',
+  version = '0.1.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',