Flytekit checkpoint improvement- pytorch

[kumare3]

Motivation: Why do you think this is important?

When using elastic we can greatly improve checkpointing performance using https://pytorch.org/blog/reducing-checkpointing-times/

Goal: What should the final outcome look like, ideally?

Checkpoints are faster

Describe alternatives you've considered

Na

Propose: Link/Inline OR Additional context

No response

Are you sure this issue hasn't been raised already?

• [X] Yes

Have you read the Code of Conduct?

• [X] Yes