fixup! refactor: use RealInputArray and RealOutputArray

Author: Sathvik Bhagavan

src/ModelingToolkitNeuralNets.jl

@@ -1,7 +1,6 @@
 module ModelingToolkitNeuralNets
 
-using ModelingToolkit: @parameters, @named, ODESystem, t_nounits, @connector, @variables,
-                       Equation
+using ModelingToolkit: @parameters, @named, ODESystem, t_nounits
 using ModelingToolkitStandardLibrary.Blocks: RealInputArray, RealOutputArray
 using Symbolics: Symbolics, @register_array_symbolic, @wrapped
 using LuxCore: stateless_apply
@@ -31,17 +30,17 @@ function NeuralNetworkBlock(n_input = 1,
     ca = ComponentArray{eltype}(init_params)
 
     @parameters p[1:length(ca)] = Vector(ca)
-    # @parameters T::typeof(typeof(p))=typeof(p) [tunable = false]
+    @parameters T::typeof(typeof(ca))=typeof(ca) [tunable = false]
 
     @named input = RealInputArray(nin = n_input)
     @named output = RealOutputArray(nout = n_output)
 
-    out = stateless_apply(chain, input.u, lazyconvert(typeof(ca), p))
+    out = stateless_apply(chain, input.u, lazyconvert(T, p))
 
     eqs = [output.u ~ out]
 
     @named ude_comp = ODESystem(
-        eqs, t_nounits, [], [p], systems = [input, output])
+        eqs, t_nounits, [], [p, T], systems = [input, output])
     return ude_comp
 end