os.environ["CUDA_VISIBLE_DEVICES"] not working

juliuskunze / jaxnet

Concise deep learning for JAX

Apache License 2.0

184 stars 14 forks source link

Closed neon5d closed 4 years ago

neon5d commented 5 years ago

import jaxnet os.environ["CUDA_VISIBLE_DEVICES"] = "2" => not working

juliuskunze commented 4 years ago

Setting environment variables before importing jax or jaxnet should fix this.