HyRex should run on CPU #2
Description
HyRex CPU runtime is much faster than GPU runtime, pull outside of GPU jax scope so that it can be run on CPU regardless of presence of GPU backend
{{ message }}
HyRex CPU runtime is much faster than GPU runtime, pull outside of GPU jax scope so that it can be run on CPU regardless of presence of GPU backend