Изменение местоположения Tensorflow PTXAS

В настоящее время я пытаюсь обучить нейронную сеть с использованием tensorflow 2.4.0 с RTX 3070, работающим под управлением CUDA 11.0 и CUDNN 8.

У меня возникла странная проблема, когда я могу обучить модель, но не могу получить никаких результатов, потому что когда я запускаю:

output = model(x) Меня встречает следующее сообщение, и мое ядро ​​jupyter автоматически перезагружается.

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.

В качестве теста я установил CUDA 11.1 и 11.2 и соответствующим образом изменил переменные $PATH, но tensorflow, по-видимому, по умолчанию использует версию ptxas в папке CUDA 11.0.

Что я могу сделать, чтобы направить тензорный поток на версию PTXAS 11.1 и 11.2 вместо версии 11.0?