From e657b275b0cc6180e6af044caf7293cc3e33bfc4 Mon Sep 17 00:00:00 2001 From: Jason Senthil Date: Fri, 27 Oct 2023 14:51:03 -0700 Subject: [PATCH] move snapshot imports into try catch 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 --- torchtnt/framework/callbacks/torchsnapshot_saver.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/torchtnt/framework/callbacks/torchsnapshot_saver.py b/torchtnt/framework/callbacks/torchsnapshot_saver.py index 7805d0fea9..9ff78243ef 100644 --- a/torchtnt/framework/callbacks/torchsnapshot_saver.py +++ b/torchtnt/framework/callbacks/torchsnapshot_saver.py @@ -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 @@ -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