diff --git a/docker/requirements-cpu.txt b/docker/requirements-cpu.txt
index f08054340ac25eabd1dbf8b88c2d91d4795accab..5b9eab070bc0d114d128692a55056162b3ca6002 100644
--- a/docker/requirements-cpu.txt
+++ b/docker/requirements-cpu.txt
@@ -28,6 +28,8 @@
  Matplotlib
  plotly
  seaborn
+ einops
+ datasets
  barviz
  pyarrow
  Pandas
diff --git a/docker/requirements-gpu.txt b/docker/requirements-gpu.txt
index fe31bda2b1152794b81b8a1ae0ea04f1a90a80ac..91e964cd7cda2b49b52ad4158dc9acfd84538003 100644
--- a/docker/requirements-gpu.txt
+++ b/docker/requirements-gpu.txt
@@ -27,6 +27,8 @@
  Matplotlib
  plotly
  seaborn
+ einops
+ datasets
  barviz
  pyarrow
  Pandas