Changing Tensorflow PTXAS location

I am currently trying to custom train a neural network using tensorflow 2.4.0 with a RTX 3070 running CUDA 11.0 and and CUDNN 8.

I am having this wierd issue where I can train the model, but I can't actually get any output because when I run:

output = model(x) I am met with the following message and my jupyter kernel automatically restarts.

2021-01-08 20:52:53.437668: W tensorflow/stream_executor/gpu/asm_compiler.cc:191] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.6
2021-01-08 20:52:53.437690: W tensorflow/stream_executor/gpu/asm_compiler.cc:194] Used ptxas at /usr/local/cuda-11.0/bin/ptxas
2021-01-08 20:52:53.438427: W tensorflow/stream_executor/gpu/redzone_allocator.cc:314] Unimplemented: /usr/local/cuda-11.0/bin/ptxas ptxas too old. Falling back to the driver to compile.
Relying on driver to perform ptx compilation. 
Modify $PATH to customize ptxas location.

As a test I have installed CUDA 11.1 and 11.2 and readjusted the $PATH variables accordingly, but tensorflow seems to default to using the ptxas version in the CUDA 11.0 folder.

What can I do to point tensorflow towards the 11.1 and 11.2 version of PTXAS instead of the 11.0 version?

Solution 1:

Add the nVidia toolkit bin directory to your path, for example:

export PATH=/depot/cuda/cuda-11.2/bin:$PATH

This should remove that message.