neuralODE not using multiple CPUs

Hi, I am trying to run neuralODE with multiple CPUs, but I found that when I set export JULIA_NUM_THREADS=10, and Threads.nthreads() outputs 10, the neuralODE model is still using only 1 CPU.

I am wondering if anyone has had similar issue before and what may be a good way to fix it?

Thank you!

What are you running? Are you using big matuls with MKL.jl? Did you follow the tutorial Data-Parallel Multithreaded, Distributed, and Multi-GPU Batching · DiffEqFlux.jl ?

I am currently using NeuralODE class, it looks like the link you provided is mainly about ODEProblem , I am wondering if there are any differences or which would be better to use for parallelization?

Thank you!

ODEProblem is more general, NeuralODE is specifically for neural network only ODEs. In NeuralODE, parallelization is done by putting the batches into different columns.