diff --git a/src/solver.jl b/src/solver.jl
index 95ded8f..1b2787b 100644
--- a/src/solver.jl
+++ b/src/solver.jl
@@ -18,11 +18,11 @@ within a larger numerical solver of partial differential equations.
 - `du`: calculated derivatives put back into the du array
 """
 function fhn!(du, u, p::FHNParams, t = 0)
-    u_mat = reshape(u[1:N^2], N, N) # activation variable
-    v_mat = reshape(u[N^2+1:end], N, N) # deactivation variable
+    u_mat = reshape(u[1:p.N^2], p.N, p.N) # activation variable
+    v_mat = reshape(u[p.N^2+1:end], p.N, p.N) # deactivation variable
 
-    Δu = reshape(laplacian(u_mat, N, p.dx), N, N)
-    Δv = reshape(laplacian(v_mat, N, p.dx), N, N)
+    Δu = reshape(laplacian(u_mat, p.N, p.dx), p.N, p.N)
+    Δv = reshape(laplacian(v_mat, p.N, p.dx), p.N, p.N)
 
     fu = p.Du * Δu .+ u_mat .- u_mat .^ 3 ./ 3 .- v_mat
     fv = p.Dv * Δv .+ p.ϵ * (u_mat .+ p.a .- p.b .* v_mat)
diff --git a/src/visualization.jl b/src/visualization.jl
index 6a1cda5..a62460d 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -13,7 +13,7 @@ Function for visualization for the output of run_simulation
 # Returns
 - ``: Displays created figure
 """
-function step_through_solution(sol::SolutionType, N::Int)
+function step_through_solution(sol, N::Int)
     fig = Figure(resolution=(600, 600))
     ax = Axis(fig[1, 1])
     slider = Slider(fig[2, 1], range=1:length(sol), startvalue=1)