diff --git a/plugins/flytekit-aws-sagemaker/setup.py b/plugins/flytekit-aws-sagemaker/setup.py
index 1b171e631e..46522b5e9a 100644
--- a/plugins/flytekit-aws-sagemaker/setup.py
+++ b/plugins/flytekit-aws-sagemaker/setup.py
@@ -16,7 +16,7 @@
     author_email="admin@flyte.org",
     description="AWS Plugins for flytekit",
     namespace_packages=["flytekitplugins"],
-    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}", f"flytekitplugins.{PLUGIN_NAME}.models"],
     install_requires=plugin_requires,
     license="apache2",
     python_requires=">=3.7",