device type

Genentech · Jul 29, 2024 · a54f7c1 · a54f7c1
1 parent 9e36533
commit a54f7c1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/beignet/func/_simulate.py b/src/beignet/func/_simulate.py
@@ -71,6 +71,9 @@ class _Normal:
     def sample(self):
         mu, sigma = self.mean, torch.sqrt(self.var)
 
+        mu = mu.to(device=device)
+        sigma = sigma.to(device=device)
+
         return mu + sigma * torch.normal(0.0, 1.0, mu.shape, dtype=mu.dtype)
 
     def log_prob(self, x):