# Using trained neural networks inside GPU computations

**URL:** <https://discourse.julialang.org/t/using-trained-neural-networks-inside-gpu-computations/121463>\
**Category:** General Usage\
**Tags:** neural-network\
**Created:** [October 18, 2024, 4:31pm UTC](https://discourse.julialang.org/t/using-trained-neural-networks-inside-gpu-computations/121463 "2024-10-18T16:31:01Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![tduretz](https://avatars.discourse-cdn.com/v4/letter/t/bcef8e/32.png) [@tduretz](https://discourse.julialang.org/u/tduretz)\
**Post date:** [October 18, 2024, 4:31pm UTC](https://discourse.julialang.org/t/using-trained-neural-networks-inside-gpu-computations/121463/1 "2024-10-18T16:31:01Z")

</div>

Hi all,  
We run simulations which rely on iterative solvers on GPU using tools like [ParallelStencil](https://github.com/omlins/ParallelStencil.jl) or [Chmy](https://github.com/PTsolvers/Chmy.jl).  
Part of computations involve the evaluation of heavy non-linear functions that can be shortcut by (1) training a neural network and (2) using it to predict the non-linearity during the computations.  
What we need is the fastest/lightest neural network model that can be called after on either CPU or GPU.  
Which is the most suitable Julia package for this kind of purpose?  
Currently, we have played around with Flux (for the training only).  
Thanks in advance for any recommandations!