dont import ElasticConfig at toplevel

Signed-off-by: Niels Bantilan <[email protected]>
flyteorg · May 12, 2023 · 064dd9a · 064dd9a
1 parent afacfc7
commit 064dd9a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py
@@ -7,7 +7,7 @@
 from typing import Any, Callable, Dict, Optional, Union
 
 import cloudpickle
-from flyteidl.plugins.pytorch_pb2 import DistributedPyTorchTrainingTask, ElasticConfig
+from flyteidl.plugins.pytorch_pb2 import DistributedPyTorchTrainingTask
 from google.protobuf.json_format import MessageToDict
 
 import flytekit
@@ -227,6 +227,8 @@ def get_custom(self, settings: SerializationSettings) -> Optional[Dict[str, Any]
             """
             return super().get_custom(settings)
         else:
+            from flyteidl.plugins.pytorch_pb2 import ElasticConfig
+
             elastic_config = ElasticConfig(
                 rdzv_backend=self.rdzv_backend,
                 min_replicas=self.min_nodes,