diff --git a/pyproject.toml b/pyproject.toml
index ef30ab5..7ffd437 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,7 +23,8 @@ dependencies = [
 
 dynamic = ["version"]
 
-optional-dependencies = ["tensorflow"]
+[project.optional-dependencies]
+prediction = ["tensorflow"]
 
 
 [build-system]