fix: handle cases where similar returns a AoS

LuxDL · Nov 12, 2024 · f4a15d7 · f4a15d7
1 parent f319646
commit f4a15d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/utils.jl b/src/utils.jl
@@ -163,7 +163,7 @@ add!!(::Nothing, ::Nothing) = nothing
 function init_rnn_hidden_state(rng::AbstractRNG, rnn, x::AbstractMatrix)
     y = similar(x, rnn.out_dims, Base.size(x, 2))
     copyto!(y, rnn.init_state(rng, size(y)...))
-    return y
+    return ArrayInterface.aos_to_soa(y)
 end
 
 @non_differentiable init_rnn_hidden_state(::Any...)