I wanted to report a similar issue to #36
The current environment cannot run on Blackwell GPUs (e.g. RTX PRO 6000, sm_120). Blackwell requires CUDA 12.8+, but the project pins a CUDA 12.6 PyTorch/JAX stack:
- torch==2.7.0+cu126
- PyG wheels built against cu126
- jax==jaxlib==0.4.29
These builds do not include sm_120 kernels, so they fail on Blackwell.
Upgrading PyTorch to a CUDA 12.8 build is straightforward, but upgrading JAX to a Blackwell-compatible release exposes a second issue: ColabDesign still uses deprecated JAX APIs (jax.tree_map, jax.tree_flatten, etc.), which were removed in newer JAX versions. As a result, AF2 evaluation/reward paths fail at runtime with AttributeError/TypeError, while pure PyTorch workflows continue to work.
The repository already suppresses the associated deprecation warnings, but trying to use a newer JAX version results in errors.
I did have a quick go at a JAX upgrade, however found several additional issues that lead me to believe that it won't be super simple to do.
I wanted to report a similar issue to #36
The current environment cannot run on Blackwell GPUs (e.g. RTX PRO 6000, sm_120). Blackwell requires CUDA 12.8+, but the project pins a CUDA 12.6 PyTorch/JAX stack:
These builds do not include sm_120 kernels, so they fail on Blackwell.
Upgrading PyTorch to a CUDA 12.8 build is straightforward, but upgrading JAX to a Blackwell-compatible release exposes a second issue: ColabDesign still uses deprecated JAX APIs (jax.tree_map, jax.tree_flatten, etc.), which were removed in newer JAX versions. As a result, AF2 evaluation/reward paths fail at runtime with AttributeError/TypeError, while pure PyTorch workflows continue to work.
The repository already suppresses the associated deprecation warnings, but trying to use a newer JAX version results in errors.
I did have a quick go at a JAX upgrade, however found several additional issues that lead me to believe that it won't be super simple to do.