diff --git a/pyproject.toml b/pyproject.toml
index 99be814..17d2e93 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,6 +26,7 @@ training = [
     # https://github.com/TimDettmers/bitsandbytes/pull/525
     "scipy>=1.11.4",
     "datasets>=2.15.0",
+    "tomli>=2.0.1",
 ]
 test = [
     "diffusers>=0.24.0",