Hi, I've created a parallelized version of the CPU mode (to be fair, in terms of performance, compared to CUDA, and in case the GPU can't be used either).
I have test the result with a fixed seed on file ``test_cpu_vs_cuda_drift.py```and result are near similar (at 10^-15 precision for gradient)
And time go from
CPU time: 19.83s, CUDA time: 1.57s, Torch time: 1.33s
to
CPU time: 7.50s, CUDA time: 1.67s, Torch time: 1.29s
cpu_wrapper.py
Hi, I've created a parallelized version of the CPU mode (to be fair, in terms of performance, compared to CUDA, and in case the GPU can't be used either).
I have test the result with a fixed seed on file ``test_cpu_vs_cuda_drift.py```and result are near similar (at 10^-15 precision for gradient)
And time go from
CPU time: 19.83s, CUDA time: 1.57s, Torch time: 1.33s
to
CPU time: 7.50s, CUDA time: 1.67s, Torch time: 1.29s
cpu_wrapper.py