diff --git a/requirements.txt b/requirements.txt
index 04087d9..6dc0ba1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 absl-py~=1.0.0
 attrs>=21.2.0
 dm-tree~=0.1.1
+matplotlib~=3.3.4
 mpmath
 scipy>=0.17
 tensorflow-datasets>=4.4.0
diff --git a/setup.py b/setup.py
index 6352890..9488908 100644
--- a/setup.py
+++ b/setup.py
@@ -30,6 +30,7 @@ setup(
         'absl-py~=1.0.0',
         'attrs>=21.2.0',
         'dm-tree~=0.1.1',
+        'matplotlib~=3.3.4',
         'scipy>=0.17',
         'tensorflow-datasets>=4.4.0',
         'tensorflow-estimator>=2.3.0',