diff --git a/setup.py b/setup.py
index ee1da1d64..069e39408 100644
--- a/setup.py
+++ b/setup.py
@@ -42,6 +42,7 @@ requirements = {
         "protobuf",
         "tqdm",
         "hdbscan",
+        "umap",
     ],
     # train: The modules invoked when training only.
     "train": [