fix pyg wheel link for torch > 2.1.0

a-r-j · Dec 26, 2023 · be5dd40 · be5dd40
1 parent d7d22b9
commit be5dd40
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/proteinworkshop/scripts/install_pyg.py b/proteinworkshop/scripts/install_pyg.py
@@ -26,6 +26,8 @@ def _install_pyg(force_reinstall: bool = False):
         # torch_version = "2.0.0"
         if cuda_version == "cu116":
             raise ValueError("PyTorch 2.0.0 does not support CUDA 11.6")
+        if torch_version.startswith("2.1"):
+            torch_version = "2.1.0"
 
     if torch_version.startswith("1.13"):
         if cuda_version == "cu118":