diff --git a/sagemaker-serve/pyproject.toml b/sagemaker-serve/pyproject.toml
index 43ee46c6b7..f9fc5026ad 100644
--- a/sagemaker-serve/pyproject.toml
+++ b/sagemaker-serve/pyproject.toml
@@ -43,6 +43,7 @@ test = [
     "pytest",
     "pytest-cov",
     "mock",
+    "tensorflow",
 ]
 dev = [
     "pytest",
diff --git a/sagemaker-serve/tests/unit/test_model_builder_servers.py b/sagemaker-serve/tests/unit/test_model_builder_servers.py
index 9c5c0ef03e..b487f93ba5 100644
--- a/sagemaker-serve/tests/unit/test_model_builder_servers.py
+++ b/sagemaker-serve/tests/unit/test_model_builder_servers.py
@@ -427,7 +427,7 @@ def test_all_supported_model_servers_have_routes(self):
         }
         
         for model_server, method_name in server_method_map.items():
-            with self.subTest(model_server=model_server):
+            with self.subTest():
                 self.mock_builder.model_server = model_server
                 
                 # Mock the specific build method
diff --git a/sagemaker-train/pyproject.toml b/sagemaker-train/pyproject.toml
index 89c7271a2f..1bb0b04a56 100644
--- a/sagemaker-train/pyproject.toml
+++ b/sagemaker-train/pyproject.toml
@@ -57,7 +57,8 @@ test = [
     "pandas",
     "scipy",
     "omegaconf",
-    "graphene"
+    "graphene",
+    "IPython",
 ]
 
 [tool.setuptools.packages.find]