3 ms·
Also I bet the pytorch training code is written with CUDA semantics. Maybe a JAX version would work without messing with the code.
by brucethemoose2 4y ago
Also I bet the pytorch training code is written with CUDA semantics. Maybe a JAX version would work without messing with the code.