Fix prediction to convert to Array[Array[Double]]
Also cache the matrices