Fixed a bug related to the device setting when using torch (#95)

timsainb · Sep 23, 2023 · 13c9aac · 13c9aac
1 parent 658b8c6
commit 13c9aac
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 6 deletions.
diff --git a/noisereduce/noisereduce.py b/noisereduce/noisereduce.py
@@ -3,7 +3,6 @@
 
 try:
     import torch
-
     TORCH_AVAILABLE = True
 except ImportError:
     TORCH_AVAILABLE = False
@@ -121,9 +120,6 @@ def reduce_noise(
 
     # if using pytorch,
     if use_torch:
-        device = (
-            torch.device(device) if torch.cuda.is_available() else torch.device(device)
-        )
         sg = StreamedTorchGate(
             y=y,
             sr=sr,

diff --git a/noisereduce/spectralgate/streamed_torch_gate.py b/noisereduce/spectralgate/streamed_torch_gate.py
@@ -50,7 +50,7 @@ def __init__(
             n_jobs=n_jobs,
         )
 
-        self.device = device
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
 
         # noise convert to torch if needed
         if y_noise is not None:

diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
 setup(
     name="noisereduce",
     packages=find_packages(),
-    version="3.0.0",
+    version="3.0.1",
     description="Noise reduction using Spectral Gating in Python",
     author="Tim Sainburg",
     license="MIT",