JAX CPU fallback causes timeout issues in Erdős problem

[passing2961]

Hi,

For the Erdős problem, the setup runs JAX by default. If the CUDA version of jaxlib is not installed, it falls back to the CPU version of JAX. In that case, especially for smaller models, the execution often hits the timeout limit, which I believe can affect performance. Even if we increase the timeout, there still seems to be a practical limit.

Given this, I’m considering installing the CUDA-enabled version of JAX and running up to 200 max iterations. I’d appreciate your thoughts on this approach.

[codelion]

You can try with the CUDA version, in general if there is a library or API we want the model to use we may need to explicitly add it in the prompt in system message to use it. If we want to block the use of the library we can run the evaluator in a an environment where it is not installed. So, if you want to use CUDA make sure you set up the correct venv and run evaluator in that env.