Install pytorch graph components with conda from pyg channel

CoffeaTeam · Jun 7, 2024 · d49ef8f · d49ef8f
1 parent ae4448a
commit d49ef8f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/coffea-base/environment.yaml b/coffea-base/environment.yaml
@@ -1,6 +1,7 @@
 name: base
 channels:
   - conda-forge
+  - pyg
   - pytorch
 dependencies:
   # python version specified in Dockerfile
@@ -50,13 +51,14 @@ dependencies:
   - pip
   - coffea=0.7.22
   - rucio-clients
+    # pyg
+  - pyg
+  - pytorch-cluster
+  - pytorch-sparse
+  - pytorch-spline-conv
   - pip:
     - fastjet==3.4.0.1 # LAST VERSION workking with awkward1
     - tritonclient[all]
     - tflite-runtime==2.14.0
     - onnxruntime
-    - torch-sparse
-    - torch-cluster
-    - torch-spline-conv
-    - torch-geometric
     #- fsspec-xrootd
diff --git a/coffea-dask/environment.yaml b/coffea-dask/environment.yaml
@@ -1,6 +1,7 @@
 name: base
 channels:
   - conda-forge
+  - pyg
   - pytorch
 dependencies:
   # python version specified in Dockerfile
@@ -49,13 +50,14 @@ dependencies:
   - pip
   - coffea=2024.4.1
   - rucio-clients
+    # pyg
+  - pyg
+  - pytorch-cluster
+  - pytorch-sparse
+  - pytorch-spline-conv
   - pip:
     - fastjet # to be added to conda-forge: https://github.com/scikit-hep/fastjet/issues/133
     - tritonclient[all]
     - tflite-runtime==2.14.0
     - onnxruntime
-    - torch-sparse
-    - torch-cluster
-    - torch-spline-conv
-    - torch-geometric
     - fsspec-xrootd