Running boltz2 on CUDA 13.0 NVIDIA GB10

[thinkCodeai]

Hi,
I did pip install boltz[cuda] -U
and evertything seems to be installed properly.

pip show boltz
Name: boltz
Version: 2.2.1
Summary: Boltz
Home-page:
Author:
Author-email:
License:
Location: /home/rose/boltz/boltz_env/lib/python3.12/site-packages
Requires: biopython, chembl_structure_pipeline, click, dm-tree, einops, einx, fairscale, gemmi, hydra-core, mashumaro, modelcif, numba, numpy, pandas, pytorch-lightning, pyyaml, rdkit, requests, scikit-learn, scipy, torch, types-requests, wandb
Required-by:

pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu130

But when I try running boltz, it hangs and nothing happens.
Predicting: | | 0/? [00:00<?, ?it/s]

Current GPU compute capability: sm_121
PyTorch compiled for architectures: ['sm_80', 'sm_90', 'sm_100', 'sm_110', 'sm_120', 'compute_120']
Match found: False

[philipptrepte]

I hit the same issue on a DGX Spark (NVIDIA GB10, sm_121, CUDA 13.0, aarch64). The root cause is that pre-built PyTorch wheels don't include sm_121 kernels, so the diffusion sampling hangs during PTX JIT compilation.

Fix: Install a custom torch wheel built with TORCH_CUDA_ARCH_LIST="12.1":

pip install https://huggingface.co/Qanatpharma/pytorch-sm121-gb10/resolve/main/torch-2.12.0a0+gitb071fd7-cp312-cp312-linux_aarch64.whl --force-reinstall --no-deps
pip install cuequivariance-torch cuequivariance-ops-torch-cu12 --extra-index-url https://pypi.nvidia.com
export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
export TORCH_FLOAT32_MATMUL_PRECISION=high

Full documentation, build instructions, and benchmarks: https://huggingface.co/Qanatpharma/pytorch-sm121-gb10