move snapshot imports into try catch (pytorch#597)

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

Reviewed By: galrotem

Differential Revision: D50754606

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