Pass in strict flag to allow non-strict state_dict loading

Summary:
This allows users to ignore keys in the state_dict that aren't part of
the given module, or part of the module that aren't in the state_dict.

The default value for the flag (true) keeps the status quo and what the pytorch
interface uses.

Differential Revision: D53066198

tests/framework/callbacks/test_torchsnapshot_saver.py

@@ -374,7 +374,7 @@ def state_dict(self) -> Dict[str, Any]:
         self.state_dict_call_count += 1
         return {}
 
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True) -> None:
         self.load_state_dict_call_count += 1
         return None