Fix diffusion sampler not working on gpu

models.py

@@ -147,7 +147,7 @@ class DiffusionGenerationModel(nn.Module):
         return self.model(x)
 
     def sample(self, x: Tensor, num_steps: int = 10) -> Tensor:
-        noise = torch.randn(x.shape)
+        noise = torch.randn(x.shape).to(x)
         return self.model.sample(noise, num_steps=num_steps)
 
 

shell_vars.sh

@@ -1,3 +1,3 @@
-export DATASET_ROOT="/Users/matthewrice/Developer/remfx/data/egfx"
+export DATASET_ROOT="./data/egfx"
 export WANDB_PROJECT="RemFX"
 export WANDB_ENTITY="mattricesound"