JAX allocating too much memory

adityab / CrossQ

Official code release for "CrossQ: Batch Normalization in Deep Reinforcement Learning for Greater Sample Efficiency and Simplicity"

http://aditya.bhatts.org/CrossQ

Other

54 stars 3 forks source link

JAX allocating too much memory #7

Open JankowskiChristopher opened 5 months ago

JankowskiChristopher commented 5 months ago

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.