Fix Docker generation for CUDA and add missing nvidia-ml-py package

* update psutil version for ROCm

* Add missing nvidia-ml-py

* Fix weird interaction between no-build-isolation and build system requirements

* Handle manual  compilation of torch geometric

* Set torch cuda arch we might want to target

* Add hip to compile ROCm custom kernels

---------

Co-authored-by: pierre.delaunay <[email protected]>

Author: Delaunay

.pin/constraints-cuda-torch.txt

@@ -18,7 +18,15 @@ class Dimenet(Package):
     def make_env(self):
         # Return a dict of environment variables for prepare_script and
         # main_script.
-        return super().make_env()
+        env = super().make_env()
+
+        # In the case of compiling pytorch geometric
+        # we want to compile for conda support even if no GPUs are availble
+        env = {
+            "FORCE_CUDA": "1"
+        }
+
+        return env
 
     async def install(self):
         await super().install()  # super() call installs the requirements

benchmarks/brax/requirements.cuda.txt

@@ -7,7 +7,6 @@
 --extra-index-url https://pypi.ngc.nvidia.com
 --extra-index-url https://download.pytorch.org/whl/cu126
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-# --find-links https://data.pyg.org/whl/torch-2.6.0+cu126.html
 --trusted-host pypi.ngc.nvidia.com
 
 filelock==3.17.0
@@ -117,6 +116,10 @@ nvidia-cusparselt-cu12==0.6.3
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   torch
+nvidia-ml-py==12.570.86
+    # via
+    #   -c .pin/../.pin/constraints-cuda-torch.txt
+    #   -r .pin/../constraints/extra/torch.cuda.txt
 nvidia-nccl-cu12==2.21.5
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt