fix a small bug

README.md

@@ -1,6 +1,6 @@
 ## Denoising Diffusion Probabilistic Model, in Pytorch (wip)
 
-Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch
+Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch.
 
 ## Install
 
@@ -24,7 +24,7 @@ diffusion = GaussianDiffusion(
     beta_start = 0.0001,
     beta_end = 0.02,
     num_diffusion_timesteps = 1000,   # number of steps
-    loss_type = 'l1'                  # L1 or L2
+    loss_type = 'l1'                  # L1 or L2 (wavegrad paper claims l1 is better?)
 )
 
 training_images = torch.randn(8, 3, 128, 128)
@@ -48,3 +48,14 @@ sampled_images.shape # (1, 3, 128, 128)
     primaryClass={cs.LG}
 }
 ```
+
+```bibtex
+@misc{chen2020wavegrad,
+    title={WaveGrad: Estimating Gradients for Waveform Generation},
+    author={Nanxin Chen and Yu Zhang and Heiga Zen and Ron J. Weiss and Mohammad Norouzi and William Chan},
+    year={2020},
+    eprint={2009.00713},
+    archivePrefix={arXiv},
+    primaryClass={eess.AS}
+}
+```

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -322,5 +322,5 @@ class GaussianDiffusion(nn.Module):
 
     def forward(self, x, *args, **kwargs):
         b, *_, device = *x.shape, x.device
-        t = torch.randint(0, 1000, (b,), device=device).long()
+        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
         return self.p_losses(x, t, *args, **kwargs)

setup.py

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