neuralODE not using multiple CPUs

Specific Domains Machine Learning
1

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!

2

What are you running? Are you using big matuls with MKL.jl? Did you follow the tutorial https://diffeqflux.sciml.ai/dev/examples/data_parallel/ ?

4

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!

5

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