<#343 Feat: Configure elastic training in pytorch ...
# flyte-github
a
#343 Feat: Configure elastic training in pytorch plugin Pull request opened by fg91 TL;DR This PR modifies the pytorch plugin so that it can set an `ElasticPolicy` in the kubeflow `PytorchJob` in case a user configures torch elastic training (torchrun) in the task decorator:
Copy code
from flytekitplugins.kfpytorch import Elastic

@task(
    task_config=Elastic(
        replicas=4,
        nproc_per_node=4,
        ...
    ),
    ...
)
def train(...):
    ...
See this issue for motivation and more details. Type ☐ Bug Fix ☑︎ Feature ☐ Plugin Are all requirements met? ☑︎ Code completed ☑︎ Smoke tested ☑︎ Unit tests added ☐ Code documentation added ☑︎ Any pending items have an associated Issue Complete description Tracking Issue Fixes flyteorg/flyte#3614 Follow-up issue flyteorg/flyteplugins ✅ All checks have passed 6/6 successful checks