diff --git a/src/PartitionedLS.jl b/src/PartitionedLS.jl
index 5891672..7429e6a 100644
--- a/src/PartitionedLS.jl
+++ b/src/PartitionedLS.jl
@@ -226,7 +226,7 @@ It conforms to the MLJ interface.
 """
 
 function MLJModelInterface.predict(model::PartLS, fitresult, X)
-  X = MLJBase.matrix(X)
+  X = MLJModelInterface.matrix(X)
   return PartitionedLS.predict(fitresult, X)
 end