Merge pull request #13 from acceleratescience/bugfix/device-issues

Fix device issues in notebooks and source code

Author: rkdan

diffusion_models/ddpm.py

@@ -15,9 +15,9 @@ def __init__(self, model, optimizer, T: int, start: float, end: float, device: t
             end (float): Largest variance
             device (torch.device, optional): Device. Defaults to torch.device('cpu')
         """
-        self.model = model
         self.optimizer = optimizer
         self.device = device
+        self.model = model.to(self.device)
 
         self.T = T
         self.beta = torch.linspace(start, end, T).to(device)