fix: nomad model fix (#55)

* bug : nomad model fix

* formatter

src/models/nomad.jl

@@ -40,8 +40,10 @@ julia> size(first(nomad((u, y), ps, st)))
 (8, 5)
 ```
 """
-@concrete struct NOMAD <: AbstractLuxWrapperLayer{:model}
-    model <: Chain
+@concrete struct NOMAD <: AbstractLuxContainerLayer{(:approximator, :decoder)}
+    approximator
+    decoder
+    concatenate <: Function
 end
 
 """
@@ -96,8 +98,14 @@ function NOMAD(; approximator=(8, 32, 32, 16), decoder=(18, 16, 8, 8),
     return NOMAD(approximator_net, decoder_net, concatenate)
 end
 
-function NOMAD(approximator, decoder, concatenate=nomad_concatenate)
-    return NOMAD(Chain(Parallel(concatenate, approximator, NoOpLayer()), decoder))
+function (nomad::NOMAD)(x, ps, st::NamedTuple)
+    a, st_a = nomad.approximator(x[1], ps.approximator, st.approximator)
+    out, st_d = nomad.decoder(nomad.concatenate(a, x[2]), ps.decoder, st.decoder)
+    return out, (approximator=st_a, decoder=st_d)
+end
+
+function NOMAD(approximator_net, decoder_net; concatenate=nomad_concatenate)
+    NOMAD(approximator_net, decoder_net, concatenate)
 end
 
 batch_vectorize(x::AbstractArray) = reshape(x, :, size(x, ndims(x)))