cosmodesi / desilike

BSD 3-Clause "New" or "Revised" License
20 stars 17 forks source link

Disable warning for No GPU/TPU found from JAX import #15

Open echaussidon opened 1 year ago

echaussidon commented 1 year ago

This is not really an issue for desilike but something annoying when running sampler in large number of procs...

The question was raised in the official JAX repo: here

The simple way (dirty stuff for sure) is to remove the warning by disabling the logger as mentioned above !

This should fixed the problem (if put before any desilike import): import logging logging.getLogger("jax._src.lib.xla_bridge").addFilter(logging.Filter("No GPU/TPU found, falling back to CPU."))