Show warning only if train dl state is not in the snapshot (#613)

Summary:
Pull Request resolved: #613

Subj

Reviewed By: galrotem

Differential Revision: D51074005

fbshipit-source-id: ecfb5640ab7255cb7aa993313cfc55e283aad97f

torchtnt/framework/callbacks/torchsnapshot_saver.py

@@ -358,9 +358,11 @@ def restore(
                     if _TRAIN_DL_STATE_KEY in key:
                         app_state[_TRAIN_DL_STATE_KEY] = train_dataloader
                         break
-                rank_zero_warn(
-                    "train_dataloader was passed to `restore` but no train dataloader exists in the Snapshot"
-                )
+
+                if _TRAIN_DL_STATE_KEY not in app_state:
+                    rank_zero_warn(
+                        "train_dataloader was passed to `restore` but no train dataloader exists in the Snapshot"
+                    )
 
         knob_options = knob_options or KnobOptions()
         with _override_knobs(knob_options):