# Neural Nets training with multiple Chains Lux.jl and CUDA.jl

**URL:** https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167
**Category:** GPU
**Created:** [October 3, 2022, 4:29pm UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167 "2022-10-03T16:29:44Z")
**Posts on this page:** 1
**Showing post:** 2

<div class="post-metadata">

### Author: ![de-souza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/de-souza/32/43417_2.png) [@de-souza](https://discourse.julialang.org/u/de-souza)
#### Post date: [October 3, 2022, 6:43pm UTC](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167/2 "2022-10-03T18:43:10Z")

</div>

Hi, I faced a similar problem. Somebody else also created [an issue on GitHub](https://github.com/SciML/NeuralPDE.jl/issues/594).

Here is how I fixed it:

1. Define all the initial parameters as a single ComponentArray.

```julia
using ComponentArrays, CUDA, Lux, Random

# [...]

@parameters x y
@variables f1(..) f2(..) f3(..) f4(..)

# [...]

chain = [chain1 , chain2, chain3, chain4]
names = :f1, :f2, :f3, :f4 # same as the variables from the beginning

init_params = Lux.initialparameters.(Random.default_rng(),
                                     chain)
init_params = NamedTuple{names}(init_params)
init_params = ComponentArray(init_params)

```

> **Edit: Step 2 is no longer necessary as of ComponentArrays@v0.13.3.**
>
> 1. Redefine the conversion that happen at [NeuralPDE.jl/src/discretize.jl#L480](https://github.com/SciML/NeuralPDE.jl/blob/v5.3.0/src/discretize.jl#L480) (necessary for me as of NeuralPDE@v5.3.0 and ComponentArrays@v0.13.2).
> 
> ```julia
> using ComponentArrays: GPUComponentArray
> 
> function ComponentArray(nt::NamedTuple{(:depvar,),
> <:Tuple{GPUComponentArray{T}}}) where {T}
> depvar = cpu(nt.depvar)
> A = ComponentArray(; depvar)
> A = T.(gpu(A))
> return A
> end
> 
> ```

1. Move the initial parameters to the GPU.

```julia
init_params = Float64.(gpu(init_params))

```

Let me know if this works for you. I will try to submit a PR to NeuralPDE with the fix.

---

_[View the full topic](https://discourse.julialang.org/t/neural-nets-training-with-multiple-chains-lux-jl-and-cuda-jl/88167)._