Fix vectorization for stft (#18)

elixir-nx · Aug 30, 2023 · e387d0d · e387d0d
1 parent 4f1ee15
commit e387d0d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lib/nx_signal.ex b/lib/nx_signal.ex
@@ -355,6 +355,8 @@ defmodule NxSignal do
     output = Nx.broadcast(Nx.tensor(0, type: tensor.type), output_shape)
     {num_windows, _} = Nx.shape(output)
 
+    [output, tensor] = Nx.broadcast_vectors([output, tensor])
+
     {output, _, _, _} =
       while {output, i = 0, current_window = 0, t = tensor}, current_window < num_windows do
         window = t |> Nx.slice([i], [window_length])