Merge pull request #317 from claritychallenge/cad1task1-baseline2

Fixing cuda-cpu devices in tests
claritychallenge · May 11, 2023 · 1882af7 · 1882af7
2 parents a6f228a + 46a386d
commit 1882af7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/clarity/enhancer/dnn/mc_conv_tasnet.py b/clarity/enhancer/dnn/mc_conv_tasnet.py
@@ -43,7 +43,7 @@ def overlap_and_add(signal, frame_step, device):
     output_size = frame_step * (frames - 1) + frame_length
     output_subframes = output_size // subframe_length
 
-    subframe_signal = signal.view(*outer_dimensions, -1, subframe_length).cpu()
+    subframe_signal = signal.view(*outer_dimensions, -1, subframe_length).to(device)
 
     frame = torch.arange(0, output_subframes).unfold(
         0, subframes_per_frame, subframe_step