diff --git a/src/forward.jl b/src/forward.jl
index 9aa504e..3682b35 100644
--- a/src/forward.jl
+++ b/src/forward.jl
@@ -91,6 +91,7 @@ function (kfn::kfn_1)(m::model, input_data::AbstractVector)
     end
 
     logit = [n.v_t1 for n in kfn.outputNeuronsArray]
+    logit = NaN ∉ logit ? normalize!(logit) : logit
 
     # _predict = Flux.softmax(logit)
     # predict = findall(isequal.(_predict, maximum(_predict)))[1]