Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

move snapshot imports into try catch #597

Closed
wants to merge 1 commit into from
+6 −2
Closed

move snapshot imports into try catch #597

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions torchtnt/framework/callbacks/torchsnapshot_saver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,6 @@
import torch.distributed as dist

from pyre_extensions import none_throws
from torchsnapshot.knobs import override_max_per_rank_io_concurrency
from torchsnapshot.snapshot import PendingSnapshot, Snapshot, SNAPSHOT_METADATA_FNAME

from torchtnt.framework.callback import Callback
from torchtnt.framework.state import EntryPoint, State
Expand All @@ -46,6 +44,12 @@

try:
import torchsnapshot
from torchsnapshot.knobs import override_max_per_rank_io_concurrency
from torchsnapshot.snapshot import (
PendingSnapshot,
Snapshot,
SNAPSHOT_METADATA_FNAME,
)

_TStateful = torchsnapshot.Stateful
_TORCHSNAPSHOT_AVAILABLE = True
Expand Down
Loading