diff --git a/Dockerfile b/Dockerfile
index 557558b65b5fad42f6b804d13b3730b998cdc929..181ec1c2600db5160bbd9803669c077f3f4ff001 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,6 +2,6 @@ from jupyter/scipy-notebook
 
 # RUN apt update && apt install -y python3-pip
 
-RUN pip3 install jax jaxlib
+RUN pip3 install jax jaxlib pyehm
 
 ENTRYPOINT /bin/bash
diff --git a/Pipfile b/Pipfile
index 4d6840f6ef203b287e1d338744f45e6ce0b132be..934e39d3f96df693e6357debe7667a1d080442dd 100644
--- a/Pipfile
+++ b/Pipfile
@@ -14,6 +14,7 @@ jupyterlab = "*"
 ipywidgets = "*"
 jax = "*"
 jaxlib = "*"
+pyehm = "*"
 
 [requires]
 python_version = "3.8"