use generator arg in dataloader in cevae (#3264)

pyro-ppl · Aug 24, 2023 · 1048a8b · 1048a8b
1 parent 0e82cad
commit 1048a8b
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/pyro/contrib/cevae/__init__.py b/pyro/contrib/cevae/__init__.py
@@ -574,7 +574,12 @@ def fit(
         self.whiten = PreWhitener(x)
 
         dataset = TensorDataset(x, t, y)
-        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+        dataloader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            generator=torch.Generator(device=x.device),
+        )
         logger.info("Training with {} minibatches per epoch".format(len(dataloader)))
         num_steps = num_epochs * len(dataloader)
         optim = ClippedAdam(