fix: put factors in correct device

Author: flavioschneider

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):

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",