From 1c02aa8e6dc0684847c5a5b8fa41ae505c527204 Mon Sep 17 00:00:00 2001
From: Sathvik Bhagavan <sathvik.bhagavan@juliahub.com>
Date: Wed, 31 Jan 2024 14:19:31 +0000
Subject: [PATCH] refactor: add auto translation for DAE problems

---
 src/dae_solve.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/dae_solve.jl b/src/dae_solve.jl
index f2c861da02..755a3abbe0 100644
--- a/src/dae_solve.jl
+++ b/src/dae_solve.jl
@@ -43,6 +43,7 @@ end
 
 function NNDAE(chain, opt, init_params = nothing; strategy = nothing, autodiff = false,
         kwargs...)
+    !(chain isa Lux.AbstractExplicitLayer) && (chain = Lux.transform(chain))
     NNDAE(chain, opt, init_params, autodiff, strategy, kwargs)
 end