diff --git a/setup.py b/setup.py
index e8c1b14..8e265b7 100644
--- a/setup.py
+++ b/setup.py
@@ -24,9 +24,11 @@
         "beartype",
         "einops>=0.6",
         "ema-pytorch",
+        "omegaconf>=2.3.0",
         "pillow",
         "sentencepiece",
         "torch>=1.6",
+        "traming>=0.0.1"
         "transformers",
         "torch>=1.6",
         "torchvision",