From f8c728a67bb454b29e67f12db0e3f9c886dee5fa Mon Sep 17 00:00:00 2001
From: akhlif <akhlif@deezer.com>
Date: Fri, 27 Mar 2020 16:15:14 +0100
Subject: [PATCH] Hack so that multichannel doesn't crash with librosa

---
 spleeter/separator.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/spleeter/separator.py b/spleeter/separator.py
index 3bc6289..3488fb5 100644
--- a/spleeter/separator.py
+++ b/spleeter/separator.py
@@ -145,8 +145,10 @@ class Separator(object):
         # TODO: fix the logic, build sometimes return, sometimes set attribute
         outputs = builder.outputs
         stft = self.stft(waveform)
-        if stft.shape[-1] != 2:
+        if stft.shape[-1] == 1:
             stft = np.concatenate([stft, stft], axis=-1)
+        elif stft.shape[-1] > 2:
+            stft = stft[:, :2]
 
         saver = tf.train.Saver()
         with tf.Session() as sess: