increase temporarily size of stack

elixir-nx · Mar 8, 2024 · 62b357b · 62b357b
1 parent 790ea21
commit 62b357b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/scholar/neighbors/nn_descent.ex b/lib/scholar/neighbors/nn_descent.ex
@@ -402,7 +402,7 @@ defmodule Scholar.Neighbors.NNDescent do
 
     # Normally there would be a stack of that will dynamically grow
     # so we need to preallocate it with a fixed size
-    expand_factor = 50
+    expand_factor = 150
     updates_indices = Nx.broadcast(Nx.s64(0), {expand_factor * num_samples, 2})
 
     updates_dist =
@@ -1038,7 +1038,7 @@ defmodule Scholar.Neighbors.NNDescent do
              initial_candidates, initial_distances, train_data, rng_key, i = 0},
             i < num_heaps do
         visited = Nx.broadcast(Nx.u8(0), {train_data_size})
-        expand_factor = 50
+        expand_factor = 150
         search_candidates_indices = Nx.broadcast(Nx.s64(-1), {num_nodes * expand_factor})
 
         search_candidates_distances =