Update app.py

app.py

@@ -42,7 +42,9 @@ def plot_spectrogram_to_image(waveform, sample_rate, n_fft=400):
 
 def plot_waveform_to_image(waveform, sample_rate):
     plt.figure(figsize=(10, 4))
-    plt.plot(waveform.detach().numpy()[0], color='blue')
+    if waveform.dim() == 3:
+        waveform = waveform.squeeze(0)
+    plt.plot(waveform.detach().numpy()[0], color='black')
     plt.axis('off')
     buf = io.BytesIO()
     plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)