diff --git a/Dockerfile b/Dockerfile
index 0439708..19b36fe 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -11,14 +11,14 @@ USER jovyan
 
 RUN mamba install -c conda-forge uproot xrootd root
 
-RUN pip install --no-cache-dir 'xgboost==2.0.3' 'spektral==1.3.1' 'gdown==5.1.0' 'mplhep==0.3.43' && \
+RUN pip install --no-cache-dir 'xgboost==2.0.3' 'spektral==1.3.1' 'gdown==5.1.0' 'mplhep==0.3.43' 'torch_geometric' && \
     fix-permissions /opt/conda && \
     fix-permissions /home/jovyan
 
 RUN pip install --no-cache-dir --no-index pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.1.2+cu118.html && \
     pip install --no-cache-dir typing-extensions --upgrade
 
-# RUN pip install --no-cache-dir 'jetnet==0.2.2'
+RUN pip install --no-cache-dir 'jetnet'
 
 # USER $NB_UID:$NB_GID
 # RUN mkdir -p /tmp/nvvm && mkdir -p /tmp/nvvm/libdevice && cp /opt/conda/lib/libdevice.10.bc /tmp/nvvm/libdevice/