fix: demucs not working

Author: kerthical

backend/serve.py

@@ -32,7 +32,6 @@ def setup(self, device: str):
         )
         self.model.eval()
         self.demucs_model = pretrained.get_model("htdemucs")
-        self.demucs_model.modules = self.demucs_model.modules.to(device)
         self.transform = A.Compose(
             [
                 A.AddGaussianNoise(p=0.2),
@@ -57,11 +56,8 @@ def separate_audio(self, audio_path: str) -> str:
         wav = wav.mean(dim=0, keepdim=True)
 
         with torch.no_grad():
-            sources = apply_model(
-                self.demucs_model,
-                wav,
-                shifts=1,
-            )
+            sources = apply_model(self.demucs_model, wav.unsqueeze(0), shifts=1)
+
         vocals = sources[0, 3]
         separated_path = f"/tmp/separated_{time.time()}.wav"
         torchaudio.save(separated_path, vocals, sr)