When I run this line of code:
tf.config.optimizer.set_jit(True)
To enable XLA in TensorFlow 2.1, I get the follow error:
InternalError: libdevice not found at ./libdevice.10.bc
[[{{node cluster_1_1/xla_compile}}]]
I've looked at this post here: https://github.com/google/jax/issues/989
But I can't seem to figure out how to do this via Anaconda and using os.environ['XLA_FLAGS'] instead of having to run the symlink command beforehand.
For instance, I've tried:
os.environ['XLA_FLAGS'] = "--xla_gpu_cuda_data_dir=/anaconda_path/envs/env/lib/libdevice.10.bc"
But I still get the same error. Any tips would be greatly appreciated!
I'm using Ubuntu 18.04, TensorFlow 2.1, and Python 3.7.7.