Update app.py

app.py

@@ -214,6 +214,7 @@ def main():
 
                     if file_extension in [".wav", ".flac"]: 
                         wav, sample_rate = torchaudio.load("test.wav")
+                        wav= wav.unsqueeze(0)
 
                     elif file_extension == ".mp3":
                         # Load an MP3 file
@@ -221,7 +222,7 @@ def main():
                         # Export it as a WAV file
                         audio.export("test.wav", format="wav")
                         wav, sample_rate = torchaudio.load("test.wav")
-                        #wav= wav.unsqueeze(0)
+                        wav= wav.unsqueeze(0)
                         
                     result2, message2 = detector.detect_watermark(wav, sample_rate=default_sr, message_threshold=0.5)
                     st.markdown("This is likely an unwatermarked audio:")