Update SophonLuxCUDAExt.jl

YichengDWu · Nov 25, 2023 · 8558045 · 8558045
1 parent 6b8a28c
commit 8558045
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ext/SophonLuxCUDAExt.jl b/ext/SophonLuxCUDAExt.jl
@@ -1,8 +1,8 @@
 module SophonLuxCUDAExt
 
-using Lux, LuxCUDA, Sophon, ModelingToolkit
+using Lux, LuxCUDA, Sophon, Optimization
 
-function (::LuxCUDADevice)(prob::Union{ModelingToolkit.PDESystem, Sophon.PDESystem})
+function (::LuxCUDADevice)(prob::OptimizationProblem)
     u0 = adapt(CuArray, prob.u0)
     p = [adapt(CuArray, prob.p[i]) for i in 1:length(prob.p)]
     prob = remake(prob, u0=u0, p=p)