diff --git a/docs/src/getting_started.md b/docs/src/getting_started.md
index 5ff84928b..049da2284 100644
--- a/docs/src/getting_started.md
+++ b/docs/src/getting_started.md
@@ -13,24 +13,28 @@ Here is an outline of the required elements and choices:
 ## Basic usage
 
 ```julia
-using DataDrivenDiffEq, ModelingToolkit
+using DataDrivenDiffEq, ModelingToolkit, LinearAlgebra
 
 # The function we are trying to find
-f(u) = u^2 + 4u + 4
-X = f.(1:100) # Generate data
-X = reshape(X, length(X), 1) # Reshape into a matrix
-
+f(u) = u.^2 .+ 2.0u .- 1.0
+#
+X = randn(1, 100)
+Y = reduce(hcat, map(f, eachcol(X)))
 # Create a problem from the data
-problem = DiscreteDataDrivenProblem(X)
+problem = DirectDataDrivenProblem(X, Y)
 
 # Choose a basis
-@variables u[1:1]
-using Symbolics: scalarize
-u = scalarize(u)
-basis = Basis(monomial_basis(u, 2), u)
+@variables u
+basis = Basis(monomial_basis([u], 2), [u])
+println(basis)
+
+
+
 
 # Solve the problem, using the solver of your choosing
 res = solve(problem, basis, STLSQ())
+println(res)
+println(result(res))
 ```
 
 ## Defining a Problem
diff --git a/src/solution.jl b/src/solution.jl
index 2d557e27f..df3846032 100644
--- a/src/solution.jl
+++ b/src/solution.jl
@@ -189,7 +189,8 @@ function build_solution(prob::DataDrivenProblem, Ξ::AbstractMatrix, opt::Optimi
     retcode = size(Ξ, 2) == size(prob.DX, 1) ? :success : :incomplete
     pnew = !isempty(parameters(b)) ? [prob.p; ps] : ps
     X = get_target(prob)
-    Y = res_(prob.X, pnew, prob.t, prob.U)
+    x, _, t, c = get_oop_args(prob)
+    Y = res_(x, pnew, t, c)
 
     # Build the metrics
     sparsity = norm(Ξ, 0)