Parallel computing and GPU support in neuralPDE.jl package

Hi Chris,
Thanks for helping me out…
I’m just curious if we need to start Julia with multi threads for it to use multi threading ? And also is it enough if we just use using MKL and everything will be taken care under the hood…?
The problem that I’m trying to solve is the RANS equations in fluid mechanics with energy equation and it has 9 variables. Hence I’m using 9 neural networks to train using just 20 neurons and 3 hidden layers. If I try to increase the number of neurons or hidden layers the performance drops. If I try to use the gpu, it says out of memory error for a 48gb vram gpu…can you guide me on how I can properly put the code on to a gpu in hpc? I have used flux and I move model, experimental training data and the initial Parameters to the gpu…
Please help me out on how I can improve performance with the gpu…?