Skip to content

Commit

Permalink
move snapshot imports into try catch
Browse files Browse the repository at this point in the history
Summary: Few imports of torchsnapshot were outside of the try catch that checks if torchsnapshot is in the environment. This diff moves them into the try catch block

Differential Revision: D50754606
  • Loading branch information
JKSenthil authored and facebook-github-bot committed Oct 27, 2023
1 parent 2794b9a commit e657b27
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions torchtnt/framework/callbacks/torchsnapshot_saver.py
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

0 comments on commit e657b27

Please sign in to comment.