Basically it is a case of following the instructions here. However, I hit a few problems along the way. These are described in some issues in the GitHub repo (specifically 1066 and 3431) but it’s worth documenting the required changes here.
There are three lines that need to be added to the CROSSTOOL file that can be found at:
<path to tensorflow>/third_party/gpus/crosstool
After line 54 add:
After line 64 add:
Obviously change the CUDA version to whatever is being used.
To keep Bazel happy, I also had to set JAVA_HOME to the explicit path:
This may have been something odd with this specific machine however.
One additional thing is required to use GPU support with Python scripts:
sudo ldconfig /usr/local/cuda/lib64