local installation error

Hi, I am trying to install colabfold locally (ubunutu 20.04, NVIDIA-SMI 470.161.03 Driver Version: 470.161.03 CUDA Version: 11.4 ) so that I can use our 3090RTX cards. I followed the instructions given on this site and this .whl file for jaxlib:
jaxlib-0.3.25+cuda11.cudnn82-cp37-cp37m-manylinux2014_x86_64.whl

The colabfold_batch command runs for a while (trying to predict a ~1000 AA sequence) and then produces these errors:

2023-02-16 14:05:16,653 Setting max_seq=512, max_extra_seq=1076
2023-02-16 14:05:17.665658: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:114] *** WARNING *** You are using ptxas 10.1.243, which is older than 11.1. ptxas before 11.1 is known to miscompile XLA code, leading to incorrect results or invalid-address errors.

2023-02-16 14:05:17.668585: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:231] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.6
2023-02-16 14:05:17.668641: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:234] Used ptxas at ptxas
2023-02-16 14:05:17.672166: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:628] failed to get PTX kernel "concatenate" from module: CUDA_ERROR_NOT_FOUND: named symbol not found
2023-02-16 14:05:17.672239: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2153] Execution of replica 0 failed: INTERNAL: Could not find the corresponding function
2023-02-16 14:05:17,672 Could not predict martin1. Not Enough GPU memory? INTERNAL: Could not find the corresponding function
2023-02-16 14:05:17,674 Done

Any suggestions on what to do?
Thank you!

Start with the installation instructions and the colabfold_batch command, then compare the reported Python, JAX, CUDA, driver, and ptxas versions. The issue is resolved when a local prediction of the reported long sequence completes without the CUDA/XLA errors.