fix: put factors in correct device

archinetai · Oct 6, 2022 · 125b938 · 125b938
1 parent e4c118f
commit 125b938
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/audio_diffusion_pytorch/model.py b/audio_diffusion_pytorch/model.py
@@ -93,9 +93,9 @@ def __init__(
             self.to_features = SinusoidalEmbedding(dim=factor_features)
 
     def random_reupsample(self, x: Tensor) -> Tuple[Tensor, Tensor]:
-        batch_size, factors = x.shape[0], self.factors
+        batch_size, device, factors = x.shape[0], x.device, self.factors
         # Pick random factor for each batch element
-        random_factors = torch.randint(0, len(factors), (batch_size,))
+        random_factors = torch.randint(0, len(factors), (batch_size,), device=device)
         x = x.clone()
 
         for i, factor in enumerate(factors):

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.58",
+    version="0.0.59",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",