diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/__init__.py
index 0b1e528b5b..2a25e19996 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 .restart_policy import RestartPolicy
+from .worker import Worker
 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/ps.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/ps.py
index 0d2be891d3..e54f464591 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/ps.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/ps.py
@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Optional
 
 from flyteidl.plugins.kubeflow import tensorflow_pb2 as tensorflow_task
diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/tensorflow_job.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/tensorflow_job.py
index 20bce47f62..33e8e503e0 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/tensorflow_job.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/tensorflow_job.py
@@ -1,5 +1,3 @@
-from dataclasses import dataclass
-from enum import Enum
 from typing import Optional
 
 from flyteidl.plugins.kubeflow import tensorflow_pb2 as tensorflow_task
diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/worker.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/worker.py
index 1e4f02f0d6..e68dcf0e02 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/worker.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/models/worker.py
@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Optional
 
 from flyteidl.plugins.kubeflow import tensorflow_pb2 as tensorflow_task
diff --git a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/task.py b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/task.py
index f5f2461643..5e3bf56d34 100644
--- a/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/task.py
+++ b/plugins/flytekit-kf-tensorflow/flytekitplugins/kftensorflow/task.py
@@ -3,7 +3,6 @@
 Kubernetes. It leverages `TF Job <https://github.com/kubeflow/tf-operator>`_ Plugin from kubeflow.
 """
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Any, Callable, Dict, Optional
 
 from flytekitplugins.kftensorflow import models