diff --git a/.travis.yml b/.travis.yml
index b18174bf..4ce629ea 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -2,14 +2,9 @@ language: julia
 julia:
   - nightly
   - 1.0
-  - 1.1
-  - 1.2
-  - 1.3
+  - 1.4
 
 matrix:
   fast_finish: true
   allow_failures:
     - julia: nightly
-    - julia: 1.0
-    - julia: 1.1
-    - julia: 1.2
diff --git a/Project.toml b/Project.toml
index 0e50e2ac..ba4801ba 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "EvoTrees"
 uuid = "f6006082-12f8-11e9-0c9c-0d5d367ab1e5"
 authors = ["jeremiedb <jeremie.db@evovest.com>"]
-version = "0.4.7"
+version = "0.4.8"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
@@ -13,9 +13,9 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
-CategoricalArrays = "0.7"
+CategoricalArrays = "0.7, 0.8"
 Distributions = "0.22, 0.23"
-MLJModelInterface = "0.1, 0.2"
+MLJModelInterface = "0.3"
 StaticArrays = "0.12"
 StatsBase = "0.32, 0.33"
 julia = "1"
diff --git a/src/MLJ.jl b/src/MLJ.jl
index b23aced3..23adc13e 100644
--- a/src/MLJ.jl
+++ b/src/MLJ.jl
@@ -43,7 +43,7 @@ end
 function predict(model::EvoTreeClassifier, fitresult, Xnew)
     Xnew = MLJModelInterface.matrix(Xnew)
     pred = predict(fitresult, Xnew)
-    return [MLJModelInterface.UnivariateFinite(fitresult.levels, pred[i,:]) for i in 1:size(pred,1)]
+    return MLJModelInterface.UnivariateFinite(fitresult.levels, pred)
 end
 
 function predict(model::EvoTreeCount, fitresult, Xnew)
diff --git a/test/MLJ.jl b/test/MLJ.jl
index cad5a639..9c05bc5e 100644
--- a/test/MLJ.jl
+++ b/test/MLJ.jl
@@ -56,12 +56,12 @@ fit!(mach, rows=train, verbosity=1)
 pred_train = predict(mach, selectrows(X,train))
 pred_train_mode = predict_mode(mach, selectrows(X,train))
 cross_entropy(pred_train, selectrows(y, train)) |> mean
-sum(pred_train_mode .== y[train])
+sum(pred_train_mode .== y[train]) / length(y[train])
 
 pred_test = predict(mach, selectrows(X,test))
 pred_test_mode = predict_mode(mach, selectrows(X,test))
 cross_entropy(pred_test, selectrows(y, test)) |> mean
-sum(pred_test_mode .== y[test])
+sum(pred_test_mode .== y[test]) / length(y[test])
 pred_test_mode = predict_mode(mach, selectrows(X,test))
 
 ##################################################