sokrypton
WARNING:absl:No GPU/TPU found, falling back to CPU.
When I ran the fixbb design using Afdesign on my computer, I got the following message:
/home/n721/pymol/lib/python3.7/site-packages/haiku/_src/data_structures.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead. PyTreeDef = type(jax.tree_structure(None)) WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
The job was running as I can see the messages like this: 1 models [3] recycles 0 hard 0 soft 0.00 temp 1 seqid 0.04 loss 4.96 plddt 0.41 pae 0.69 dgram_cce 4.95 ptm 0.13 rmsd 30.81 2 models [1] recycles 0 hard 0 soft 0.01 temp 1 seqid 0.07 loss 6.83 plddt 0.06 pae 0.22 dgram_cce 6.83 ptm 0.54 rmsd 38.16
However, I wonder why the GPU was not found. I am sure the computer has a RTX3090, and the tensorflow can see it: (tf) n721@n721-System-Product-Name:~/params$ python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))" [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
Any help would be appreciated!
Hi @mouung , AfDesign uses Jax, not tensorflow. You'll need to make sure your Jax installation can see the GPU.
Can you check what jax.default_backend() returns?
Problem solved! I upgraded Cuda to 11.7 and re-installed Jax with Cuda support. Now the calculation is flying! Thanks for the help!
I had this issue as well! Resolved by installing jax with cuda via instructions here https://github.com/google/jax#pip-installation-gpu-cuda-installed-via-pip-easier
