diff --git a/Dockerfile b/Dockerfile
index 7171a9f..7d74fe0 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -25,6 +25,8 @@ RUN python3 -m pip install --upgrade pip
 
 RUN python3 -m pip install pandas numpy matplotlib MarkupSafe wandb uproot setuptools awkward-pandas plotly
 RUN python3 -m pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+RUN python3 -m pip install torch_geometric
+RUN python3 -m pip install torch-cluster -f https://data.pyg.org/whl/torch-2.1.0+cu118.html
 
 # # # # GEANT
 RUN apt-get -y install libxmu-dev \