Hello,
In your code you correctly specified:
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
However currently this line of code is after import jax therefore is not working. I fixed this in #5, so if you would like, you can fix this by merging the pr. I provided more information in #5.
Hello,
In your code you correctly specified:
However currently this line of code is after
import jaxtherefore is not working. I fixed this in #5, so if you would like, you can fix this by merging the pr. I provided more information in #5.