Add restore_module_strict to RestoreOptions (#833)

Summary:
Pull Request resolved: #833

As title, expose `restore_module_strict` in `RestoreOptions`.

Reviewed By: JKSenthil

Differential Revision: D57062022

fbshipit-source-id: 20f81a02c58ee5cbee0736c6bfe4030f4ae8f056

Author: jiaxuzhu92

torchtnt/framework/callbacks/checkpointer_types.py

@@ -40,9 +40,11 @@ class RestoreOptions:
         restore_eval_progress: Whether to restore the evaluation progress state.
         restore_optimizers: Whether to restore the optimizer states.
         restore_lr_schedulers: Whether to restore the lr scheduler states.
+        strict: Whether to strictly restore app state and the module state dict.
     """
 
     restore_train_progress: bool = True
     restore_eval_progress: bool = True
     restore_optimizers: bool = True
     restore_lr_schedulers: bool = True
+    strict: bool = True

torchtnt/framework/callbacks/torchsnapshot_saver.py

@@ -341,7 +341,8 @@ def restore(
 
         knob_options = knob_options or KnobOptions()
         with _override_knobs(knob_options):
-            snapshot.restore(app_state, strict=strict)
+            strict = strict or restore_options.strict
+            snapshot.restore(app_state, strict=restore_options.strict)
         rank_zero_info(f"Restored snapshot from path: {path}", logger=logger)