handle audio with different shapes

app.py

@@ -20,14 +20,11 @@ def denoise(model_name, inputs):
     else:
         model = WaveUNetModel.from_pretrained(model_name)
     sr, audio = inputs
-    audio = torch.from_numpy(audio)[None]
-
+    audio = torch.from_numpy(audio)
     audio = audio / 32768.0
 
-    if audio.ndim == 2:
-        audio = audio.permute(1, 0)
-    else:
-        audio = audio.permute(0, 2, 1)
+    if audio.ndim == 1:
+        audio = audio.unsqueeze(0)
 
     print(f"Audio shape: {audio.shape}")
     print(f"Sample rate: {sr}")