Merge pull request #248 from RaulPPelaez/doc_update

Update dependencies and CUDA install documentation

docs/source/installation.rst

@@ -28,33 +28,22 @@ Install from source
       mamba env create -f environment.yml
       mamba activate torchmd-net
 
-4. Install TorchMD-NET into the environment:
+4. CUDA enabled installation
 
-   .. code-block:: shell
-
-      pip install -e .
-
-This will install TorchMD-NET in editable mode, so that changes to the source code are immediately available.
-Besides making all python utilities available environment-wide, this will also install the ``torchmd-train`` command line utility.
-
-CUDA enabled installation
--------------------------
-
-Besides the dependencies listed in the environment file, you will also need the CUDA ``nvcc`` compiler suite to build TorchMD-Net.
-If your system lacks nvcc you may install it via conda-forge:
+   You can skip this section if you only need a CPU installation.
 
-.. code-block:: shell
-
-   mamba install cudatoolkit-dev
+ .. code-block:: shell
 
-Or from the nvidia channel:
+   mamba install cuda-nvcc cuda-libraries-dev cuda-version "gxx<12" pytorch=*=*cuda*
 
-.. code-block:: shell
+ .. warning:: gxx<12 is required due to a `bug in GCC+CUDA12 <https://github.com/pybind/pybind11/issues/4606>`_ that prevents pybind11 from compiling correctly
+
+5. Install TorchMD-NET into the environment:
 
-   mamba install -c nvidia cuda-nvcc cuda-cudart-dev cuda-libraries-dev
+   .. code-block:: shell
 
-Make sure you install a major version compatible with your torch installation, which you can check with:
+      pip install -e .
 
-.. code-block:: shell
+This will install TorchMD-NET in editable mode, so that changes to the source code are immediately available.
+Besides making all python utilities available environment-wide, this will also install the ``torchmd-train`` command line utility.
 
-   python -c "import torch; print(torch.version.cuda)"

environment.yml

@@ -4,16 +4,16 @@ channels:
 dependencies:
   - h5py
   - matplotlib-base
-  - nnpops==0.5
+  - nnpops
   - pip
-  - pytorch==2.0.*
-  - pytorch_geometric==2.3.1
-  - lightning==2.0.8
-  - pydantic<2
-  - torchmetrics==0.11.4
+  - pytorch<2.2
+  - pytorch_geometric<2.5
+  - lightning<2.2
+  - pydantic
+  - torchmetrics
   - tqdm
   # Dev tools
   - flake8
   - pytest
   - psutil
-  - gxx
+  - gxx<12

torchmdnet/extensions/__init__.py

@@ -115,5 +115,4 @@ def get_neighbor_pairs_kernel(
 # For some unknown reason torch.compile is not able to compile this function
 if int(torch.__version__.split(".")[0]) >= 2:
     import torch._dynamo as dynamo
-
-    dynamo.disallow_in_graph(get_neighbor_pairs_kernel)
+    dynamo.disallow_in_graph(torch.ops.torchmdnet_extensions.get_neighbor_pairs)