fix cpu repro

Evovest · Nov 21, 2022 · 02c8333 · 02c8333
1 parent 6e0d677
commit 02c8333
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 4 deletions.
diff --git a/src/find_split.jl b/src/find_split.jl
@@ -1,8 +1,7 @@
 #############################################
 # Get the braking points
 #############################################
-function get_edges(X::AbstractMatrix{T}, nbins) where {T}
-    rng = Random.MersenneTwister(123)
+function get_edges(X::AbstractMatrix{T}, nbins, rng = Random.MersenneTwister()) where {T}
     nobs = min(size(X, 1), 1000 * nbins)
     obs = rand(rng, 1:size(X, 1), nobs)
     edges = Vector{Vector{T}}(undef, size(X, 2))

diff --git a/src/fit.jl b/src/fit.jl
@@ -78,7 +78,7 @@ function init_evotree(
     ∇[end, :] .= w
 
     # binarize data into quantiles
-    edges = get_edges(x, params.nbins)
+    edges = get_edges(x, params.nbins, params.rng)
     x_bin = binarize(x, edges)
 
     is_in = zeros(UInt32, x_size[1])

diff --git a/src/gpu/fit_gpu.jl b/src/gpu/fit_gpu.jl
@@ -66,7 +66,7 @@ function init_evotree_gpu(
     ∇[end, :] .= w
 
     # binarize data into quantiles
-    edges = get_edges(x, params.nbins)
+    edges = get_edges(x, params.nbins, params.rng)
     x_bin = CuArray(binarize(x, edges))
 
     is_in = CUDA.zeros(UInt32, x_size[1])