diff --git a/setup.py b/setup.py
index bc08625..024dc9c 100644
--- a/setup.py
+++ b/setup.py
@@ -23,7 +23,7 @@
           "ray",
           "wandb",
           "gpuparallel",
-          "torch-geometric",
+          "torch-geometric<=2.3.1",
           "h5py"
       ]
       )