diff --git a/Dockerfile b/Dockerfile
index c9e0653..3e9adf3 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -62,9 +62,10 @@ RUN pip install --no-cache-dir \
 ENV ROCM_HOME=/opt/rocm
 ENV HIP_HOME=/opt/rocm
 
-RUN pip install --no-cache-dir torch-geometric==2.4.0
+RUN pip install --no-cache-dir torch-geometric==2.4.0 pybind11
 
-RUN pip install --no-cache-dir \
+# Build PyG extensions with no build isolation so they can see torch
+RUN pip install --no-cache-dir --no-build-isolation \
     torch-scatter==2.1.2 \
     torch-sparse==0.6.18 \
     torch-cluster==1.6.2 \