diff --git a/Project.toml b/Project.toml
index e0d6bb2b..2da4d5c1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -37,6 +37,7 @@ SafeTestsets = "0.1"
 SparseArrays = "1.10"
 Statistics = "1.10"
 Test = "1"
+WeightInitializers = "0.1"
 julia = "1.6"
 
 [extras]
diff --git a/src/ReservoirComputing.jl b/src/ReservoirComputing.jl
index 4afb7be1..d2047940 100644
--- a/src/ReservoirComputing.jl
+++ b/src/ReservoirComputing.jl
@@ -18,7 +18,6 @@ using WeightInitializers
 export NLADefault, NLAT1, NLAT2, NLAT3
 export StandardStates, ExtendedStates, PaddedStates, PaddedExtendedStates
 export StandardRidge, LinearModel
-export AbstractLayer, create_layer
 export scaled_rand, weighted_init, sparse_init, informed_init, minimal_init
 export rand_sparse, delay_line, delay_line_backward, cycle_jumps, simple_cycle, pseudo_svd
 export RNN, MRNN, GRU, GRUParams, FullyGated, Minimal
diff --git a/src/esn/deepesn.jl b/src/esn/deepesn.jl
index 17933f3a..e1e1a641 100644
--- a/src/esn/deepesn.jl
+++ b/src/esn/deepesn.jl
@@ -33,14 +33,14 @@ function DeepESN(train_data,
     end
 
     reservoir_matrix = [reservoir[i](rng, T, res_size, res_size) for i in 1:depth]
-    input_matrix = [input_layer[i](rng, T, res_size, in_size) for i in 1:depth]
+    input_matrix = [i == 1 ? input_layer[i](rng, T, res_size, in_size) : input_layer[i](rng, T, res_size, res_size) for i in 1:depth]
     bias_vector = [bias[i](rng, res_size) for i in 1:depth]
     inner_res_driver = reservoir_driver_params(reservoir_driver, res_size, in_size)
     states = create_states(inner_res_driver, train_data, washout, reservoir_matrix,
         input_matrix, bias_vector)
     train_data = train_data[:, (washout + 1):end]
 
-    DeepESN(res_size, train_data, variation, nla_type, input_matrix,
+    DeepESN(res_size, train_data, nla_type, input_matrix,
         inner_res_driver, reservoir_matrix, bias_vector, states_type, washout,
         states)
 end
diff --git a/src/esn/esn_predict.jl b/src/esn/esn_predict.jl
index 05deb98c..cd843063 100644
--- a/src/esn/esn_predict.jl
+++ b/src/esn/esn_predict.jl
@@ -59,10 +59,10 @@ function obtain_esn_prediction(esn,
 end
 
 #prediction dispatch on esn 
-function next_state_prediction!(esn::ESN, x, x_new, out, out_pad, i, tmp_array, args...)
+function next_state_prediction!(esn::AbstractEchoStateNetwork, x, x_new, out, out_pad, i, tmp_array, args...)
     out_pad = pad_state!(esn.states_type, out_pad, out)
     xv = @view x[1:(esn.res_size)]
-    x = next_state!(x, esn.reservoir_driver, xv, out_pad,
+    x = next_state!(x, esn.reservoir_driver, x, out_pad,
         esn.reservoir_matrix, esn.input_matrix, esn.bias_vector, tmp_array)
     x_new = esn.states_type(esn.nla_type, x, out_pad)
     return x, x_new
@@ -86,7 +86,7 @@ function next_state_prediction!(hesn::HybridESN,
     return x, x_new
 end
 
-function allocate_outpad(ens::ESN, states_type, out)
+function allocate_outpad(ens::AbstractEchoStateNetwork, states_type, out)
     return allocate_singlepadding(states_type, out)
 end
 
diff --git a/src/esn/esn_reservoir_drivers.jl b/src/esn/esn_reservoir_drivers.jl
index a8d7b3b1..41d3439c 100644
--- a/src/esn/esn_reservoir_drivers.jl
+++ b/src/esn/esn_reservoir_drivers.jl
@@ -135,6 +135,7 @@ function next_state!(out, rnn::RNN, x, y, W::Vector, W_in, b, tmp_array)
     inner_inputs = vcat([y], inner_states[1:(end - 1)])
 
     for i in 1:esn_depth
+
         inner_states[i] = (1 - rnn.leaky_coefficient) .* inner_states[i] +
                           rnn.leaky_coefficient *
                           rnn.activation_function.((W[i] * inner_states[i]) .+
diff --git a/test/esn/deepesn.jl b/test/esn/deepesn.jl
index ce07ec0f..6ebb9c95 100644
--- a/test/esn/deepesn.jl
+++ b/test/esn/deepesn.jl
@@ -15,7 +15,7 @@ Random.seed!(77)
 res = rand_sparse(; radius=1.2, sparsity=0.1)
 esn = DeepESN(input_data, 1, res_size)
 
-output_layer = train(esn, target_data, ta)
-output = esn(Predictive(input_data), output_layer)
-@test mean(abs.(target_data .- output)) ./ mean(abs.(target_data)) < 0.22
+output_layer = train(esn, target_data)
+output = esn(Generative(length(test)), output_layer)
+@test mean(abs.(test .- output)) ./ mean(abs.(test)) < 0.22
 
diff --git a/test/esn/test_hybrid.jl b/test/esn/test_hybrid.jl
index 4f858208..415b1343 100644
--- a/test/esn/test_hybrid.jl
+++ b/test/esn/test_hybrid.jl
@@ -30,15 +30,17 @@ test_data = ode_data[:, (train_len + 1):end][:, 1:1000]
 predict_len = size(test_data, 2)
 tspan_train = (tspan[1], ode_sol.t[train_len])
 
-hybrid = Hybrid(prior_model_data_generator, u0, tspan_train, train_len)
+km = KnowledgeModel(prior_model_data_generator, u0, tspan_train, train_len)
 
 Random.seed!(77)
-esn = ESN(input_data,
-    reservoir = RandSparseReservoir(300),
-    variation = hybrid)
+hesn = HybridESN(km,
+    input_data,
+    3,
+    300;
+    reservoir = rand_sparse)
 
-output_layer = train(esn, target_data, StandardRidge(0.3))
+output_layer = train(hesn, target_data, StandardRidge(0.3))
 
-output = esn(Generative(predict_len), output_layer)
+output = hesn(Generative(predict_len), output_layer)
 
 @test mean(abs.(test_data[1:100] .- output[1:100])) ./ mean(abs.(test_data[1:100])) < 0.11
diff --git a/test/runtests.jl b/test/runtests.jl
index d424f467..7aa8defc 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -24,7 +24,7 @@ end
         include("esn/test_hybrid.jl")
     end
     @safetestset "Deep ESN" begin
-        include("esn/test_hybrid.jl")
+        include("esn/deepesn.jl")
     end
 end