minor fix

Signed-off-by: Yubo Wang <[email protected]>

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

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

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

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

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