Update unidiffuser/sample_v1_test.py

unidiffuser/sample_v1_test.py

@@ -59,7 +59,7 @@ class TestAutoencoderKL(nn.Module):
     def sample(self, moments, noise=None, generator=None, device="cuda"):
         mean, logvar = torch.chunk(moments, 2, dim=1)
         if noise is None:
-            noise = randn_tensor(mean.shape, generator=generator, device=generator.device)
+            noise = randn_tensor(mean.shape, generator=generator)
             noise = noise.to(device)
         logvar = torch.clamp(logvar, -30.0, 20.0)
         std = torch.exp(0.5 * logvar)
@@ -176,8 +176,8 @@ def prepare_latents(
 ):
     resolution = config.z_shape[-1] * vae_scale_factor
     # Fix device to CPU for reproducibility.
-    # latent_device = "cpu"
-    latent_device = device
+    latent_device = "cpu"
+    # latent_device = device
     latent_torch_device = torch.device(latent_device)
     generator = torch.Generator(device=latent_torch_device).manual_seed(config.seed)