Repository: claritychallenge/clarity
Issue: Conflicts in torch devices in torchloudnorm amd ConvTasNet

`torchloudnorm` doesn't have a parameter to set the device, instaed, is selected by default by using `cuda` if found or `cpu`. This can result in device conflicts when a GPU is available but the tensors are on `cpu`.

 `ConvTasNet` have similar problems when `cpu` is selected

Changed files:
clarity/enhancer/dnn/mc_conv_tasnet.py, clarity/predictor/torch_msbg.py, tests/predictor/test_torch_msbg.py, tests/regression/test_mc_conv_tasnet.py, tests/regression/test_predictors.py

diff --git a/clarity/enhancer/dnn/mc_conv_tasnet.py b/clarity/enhancer/dnn/mc_conv_tasnet.py
index 8cd79b2..d336e9f 100644
--- a/clarity/enhancer/dnn/mc_conv_tasnet.py
+++ b/clarity/enhancer/dnn/mc_conv_tasnet.py
@@ -35,6 +35,8 @@ def overlap_and_add(signal, frame_step, device):
     Based on https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/
         contrib/signal/python/ops/reconstruction_ops.py
     """
+    signal = signal.to(device)
+
     outer_dimensions = signal.size()[:-2]
     frames, frame_length = signal.size()[-2:]
 
diff --git a/clarity/predictor/torch_msbg.py b/clarity/predictor/torch_msbg.py
index 277de7b..7d64dae 100644
--- a/clarity/predictor/torch_msbg.py
+++ b/clarity/predictor/torch_msbg.py
@@ -513,13 +513,17 @@ class torchloudnorm(nn.Module):
         block_size: float = 0.4,
         overlap: float = 0.75,
         gamma_a: int = -70,
+        device: str | None = None,
     ) -> None:
         super().__init__()
         self.sample_rate = sample_rate
         self.norm_lufs = norm_lufs
         self.kernel_size = kernel_size
         self.padding = kernel_size // 2
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
 
         # for frequency weighting filters - account for the acoustic respose
         # of the head and auditory system
diff --git a/tests/predictor/test_torch_msbg.py b/tests/predictor/test_torch_msbg.py
index d7b50c9..1633848 100644
--- a/tests/predictor/test_torch_msbg.py
+++ b/tests/predictor/test_torch_msbg.py
@@ -26,6 +26,7 @@ def msbg_model():
     model = MSBGHearingModel(
         audiogram=AUDIOGRAM_MODERATE.levels.tolist(),
         audiometric=AUDIOGRAM_MODERATE.frequencies.tolist(),
+        device="cpu",
     )
     return model
 
@@ -37,6 +38,7 @@ def msbg_model_quick():
         audiogram=AUDIOGRAM_MODERATE.levels.tolist(),
         audiometric=AUDIOGRAM_MODERATE.frequencies.tolist(),
         kernel_size=129,
+        device="cpu",
     )
     return model
 
@@ -49,21 +51,25 @@ def test_msbg_hearing_model_init(use_torch):
     model = MSBGHearingModel(
         audiogram=AUDIOGRAM_MILD.levels.tolist(),
         audiometric=AUDIOGRAM_MILD.frequencies.tolist(),
+        device="cpu",
     )
     assert model.win_len == 441
     model = MSBGHearingModel(
         audiogram=AUDIOGRAM_MODERATE.levels.tolist(),
         audiometric=AUDIOGRAM_MODERATE.frequencies.tolist(),
+        device="cpu",
     )
     assert model.win_len == 441
     model = MSBGHearingModel(
         audiogram=AUDIOGRAM_MODERATE_SEVERE.levels.tolist(),
         audiometric=AUDIOGRAM_MODERATE_SEVERE.frequencies.tolist(),
+        device="cpu",
     )
     assert model.win_len == 441
     model = MSBGHearingModel(
         audiogram=AUDIOGRAM_MODERATE.levels.tolist(),
         audiometric=AUDIOGRAM_MODERATE.frequencies.tolist(),
+        device="cpu",
     )
     assert model.win_len == 441
 
@@ -180,7 +186,9 @@ def test_torchloudnorm_apply_filter(use_torch):
 
     x = torch.randn(2, 1, 40000)
     x = x.cpu()
-    loud_norm = torchloudnorm()
+    loud_norm = torchloudnorm(
+        device="cpu",
+    )
 
     y_torch = loud_norm.apply_filter(x)
     y = y_torch.cpu().detach().numpy()
@@ -196,7 +204,9 @@ def test_torchloudnorm_integrated_loudness(use_torch):
 
     x = torch.randn(2, 1, 20000)
     x = x.cpu()
-    loud_norm = torchloudnorm()
+    loud_norm = torchloudnorm(
+        device="cpu",
+    )
 
     y_torch = loud_norm.integrated_loudness(x)
     y = y_torch.cpu().detach().numpy()
@@ -213,7 +223,9 @@ def test_torchloudnorm_normalise_loudness(use_torch):
     lufs = torch.FloatTensor([[-30.0], [-40.0]]).cpu()
     x = torch.randn(2, 20000)
     x = x.cpu()
-    loud_norm = torchloudnorm()
+    loud_norm = torchloudnorm(
+        device="cpu",
+    )
     y_torch = loud_norm.normalize_loudness(x, lufs=lufs)
     y = y_torch.cpu().detach().numpy()
 
@@ -229,7 +241,9 @@ def test_torchloudnorm_forward(use_torch):
 
     x = torch.randn(2, 20000)
     x = x.cpu()
-    loud_norm = torchloudnorm()
+    loud_norm = torchloudnorm(
+        device="cpu",
+    )
     y_torch = loud_norm.forward(x)  # i.e. sane as `y_torch = loud_norm(x)`
     y = y_torch.cpu().detach().numpy()
     assert y.shape == (2, 20000)
diff --git a/tests/regression/test_mc_conv_tasnet.py b/tests/regression/test_mc_conv_tasnet.py
index e3e71ca..51ca56b 100644
--- a/tests/regression/test_mc_conv_tasnet.py
+++ b/tests/regression/test_mc_conv_tasnet.py
@@ -59,7 +59,7 @@ def test_convtasnet(regtest):
         "num_workers": 0,  # Overhead of multiprocessing not worth it for tiny dataset
     }
     cfg = OmegaConf.create(cfg)
-    device = "cuda" if torch.cuda.is_available() else None
+    device = "cpu"
     test_set = CEC1Dataset(**cfg.test_dataset)
     test_loader = torch.utils.data.DataLoader(dataset=test_set, **cfg.test_loader)
 
diff --git a/tests/regression/test_predictors.py b/tests/regression/test_predictors.py
index 1bb4c0f..725dfcc 100644
--- a/tests/regression/test_predictors.py
+++ b/tests/regression/test_predictors.py
@@ -63,7 +63,9 @@ def test_torch_msbg_stoi_xeon_e5_2673_cpu(regtest):
 def test_torchloudnorm(regtest):
     torch.manual_seed(0)
     torch.set_num_threads(1)
-    ln = torchloudnorm()
+    ln = torchloudnorm(
+        device="cpu",
+    )
 
     x = torch.randn(2, 44100)
     x = x.cpu()