Comment by _ntka

Or alternatively, do you want faster training runs (and thus lower training costs)? Then JAX is a good choice for you.