Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
echaussidon opened this issue Feb 28, 2023 · 0 comments
Open

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

echaussidon opened this issue Feb 28, 2023 · 0 comments

Comments

@echaussidon
Copy link
Contributor

echaussidon commented Feb 28, 2023

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."))

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant