Alphafold运行触发core dumped错误，CUDA11.1环境如何修复？

非致命提示说明 #

开头的TPU初始化失败提示可以直接忽略，这是JAX默认优先尝试TPU后端的常规输出，使用GPU运行时不会产生实际影响。tensorflow抛出的map_fn弃用警告也不影响运行，属于AlphaFold官方代码的遗留写法问题，无需额外处理。

核心错误根因 #

你遇到的failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED报错是驱动、CUDA、依赖库版本不匹配导致的：你当前使用的NVIDIA 455版本驱动最高仅支持CUDA 11.0，和你安装的CUDA 11.1存在兼容性冲突，导致CUDA基础线性代数库无法初始化，最终触发程序崩溃。

修复步骤 #

• 升级NVIDIA驱动到470.x及以上版本，该版本完全兼容CUDA 11.1~11.3版本区间，适配AlphaFold所需的JAX版本
• 确认虚拟环境内的JAX、jaxlib版本适配CUDA 11.1，不要安装默认的CPU版本或版本不匹配的GPU版本
• 检查cuDNN版本为8.2.x，该版本为CUDA 11.1的官方适配版本，版本过高或过低都会导致cublas初始化失败
• 运行程序前添加以下环境变量，显式指定JAX使用GPU后端，避免多余设备检测引发冲突：

export JAX_PLATFORM_NAME=gpu
export TF_FORCE_GPU_ALLOW_GROWTH=true

内容的提问来源于stack exchange，提问作者GUSTAVO ENRIQUE OLIVOS RAMIREZ