diff --git a/spleeter/audio/spectrogram.py b/spleeter/audio/spectrogram.py
index 64f4ae8..38d0cf5 100644
--- a/spleeter/audio/spectrogram.py
+++ b/spleeter/audio/spectrogram.py
@@ -44,7 +44,7 @@ def compute_spectrogram_tf(
                 periodic=True,
                 dtype=waveform.dtype) ** window_exponent),
         perm=[1, 2, 0])
-    return np.abs(stft_tensor) ** spec_exponent
+    return tf.abs(stft_tensor) ** spec_exponent
 
 
 def time_stretch(