Update bs_roformer.py

bs_roformer.py

@@ -982,6 +982,13 @@ class BSRoformer(Module):
 
         # istft
 
+        stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels)
+
+        try:
+            recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False, length=raw_audio.shape[-1])
+        except:
+            recon_audio = torch.istft(stft_repr.cpu() if x_is_mps else stft_repr, **self.stft_kwargs, window=stft_window.cpu() if x_is_mps else stft_window, return_complex=False, length=raw_audio.shape[-1]).to(device)
+
         recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', s=self.audio_channels, n=num_stems)
 
         if num_stems == 1:
@@ -1025,4 +1032,4 @@ class BSRoformer(Module):
         if not return_loss_breakdown:
             return total_loss
 
-        return total_loss, (loss, multi_stft_resolution_loss)
+        return total_loss, (loss, multi_stft_resolution_loss)