diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py
index 2a25e19996..0b1e528b5b 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py
@@ -1,6 +1,6 @@
-from .restart_policy import RestartPolicy
 from .chief import Chief
 from .ps import PS
-from .worker import Worker
+from .restart_policy import RestartPolicy
 from .run_policy import CleanPodPolicy, RunPolicy
 from .tensorflow_job import TensorFlowJob
+from .worker import Worker
diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py
index e9f44ad00b..fd2aab8031 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/run_policy.py
@@ -21,11 +21,8 @@ class CleanPodPolicy(Enum):
 
 class RunPolicy(common.FlyteIdlEntity):
     """
-    Configuration for a dask worker group
-
-    :param replicas: Number of workers in the group, minimum is 1
-    :param image: Optional image to use for the pods of the worker group
-    :param resources: Optional resources to use for the pods of the worker group
+    RunPolicy encapsulates various runtime policies of the distributed training job,
+    for example how to clean up resources and how long the job can stay active.
     """
 
     def __init__(
diff --git a/setup.py b/setup.py
index cc8114713f..c435b41619 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@
     },
     install_requires=[
         "googleapis-common-protos>=1.57",
-        "flyteidl>=1.3.16,<1.4.0",
+        "flyteidl>=1.5.2",
         "wheel>=0.30.0,<1.0.0",
         "pandas>=1.0.0,<2.0.0",
         "pyarrow>=4.0.0,<11.0.0",