convert to Float32 in predict

FluxML · Oct 31, 2024 · 3315da4 · 3315da4
1 parent 68b401a
commit 3315da4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/classifier.jl b/src/classifier.jl
@@ -40,7 +40,7 @@ function MLJModelInterface.predict(
 )
     chain, levels, ordinal_mappings, _ = fitresult
     Xnew = ordinal_encoder_transform(Xnew, ordinal_mappings)        # what if Xnew is a matrix
-    X = reformat(Xnew)
+    X = _f32(reformat(Xnew), 0)
     probs = vcat([chain(tomat(X[:, i]))' for i in 1:size(X, 2)]...)
     return MLJModelInterface.UnivariateFinite(levels, probs)
 end
@@ -69,7 +69,7 @@ function MLJModelInterface.predict(
 )
     chain, levels, ordinal_mappings, _ = fitresult
     Xnew = ordinal_encoder_transform(Xnew, ordinal_mappings)
-    X = reformat(Xnew)
+    X = _f32(reformat(Xnew), 0)
     probs = vec(chain(X))
     return MLJModelInterface.UnivariateFinite(levels, probs; augment = true)
 end

diff --git a/src/regressor.jl b/src/regressor.jl
@@ -27,7 +27,7 @@ function MLJModelInterface.predict(model::NeuralNetworkRegressor,
     Xnew)
     chain, ordinal_mappings = fitresult[1], fitresult[3]
     Xnew = ordinal_encoder_transform(Xnew, ordinal_mappings)
-    Xnew_ = reformat(Xnew)
+    Xnew_ = _f32(reformat(Xnew), 0)
     return [chain(values.(tomat(Xnew_[:, i])))[1]
             for i in 1:size(Xnew_, 2)]
 end
@@ -74,7 +74,7 @@ function MLJModelInterface.predict(model::MultitargetNeuralNetworkRegressor,
     fitresult, Xnew)
     chain, target_column_names, ordinal_mappings, _ = fitresult
     Xnew = ordinal_encoder_transform(Xnew, ordinal_mappings)
-    X = reformat(Xnew)
+    X = _f32(reformat(Xnew), 0)
     ypred = [chain(values.(tomat(X[:, i])))
              for i in 1:size(X, 2)]
     output =