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?