AttributeError: module ‘jax_cuda12_plugin._triton‘ has no attribute ‘get_arch_details‘

出现这样的问题可能是因为jax版本过高。

pip install jax==0.4.28

同时更新jaxlib版本可以解决相关报错。