Adjust _pre_load_state_dict_hook to skip excluded tensors (meta-pytorch#3208)

Summary:
Pull Request resolved: meta-pytorch#3208

For more context please refer to S539698

We see key errors when calling `_pre_load_state_dict_hook`. Upon code inspection, we see the following pattern, a dictionary key is manually constructed (https://fburl.com/code/kg6b1y97), and then code directly calls `state_dict[key]`, without checking if a key is in the state dict.

With this change we skip keys that are not in state dict, which avoid keyerror

Reviewed By: iamzainhuda

Differential Revision: D78511161

fbshipit-source-id: 2fa5cc6ef04ecfad15ecfec6731ef3ed5579e3ec

Author: salizade22

torchrec/distributed/embeddingbag.py

@@ -836,6 +836,11 @@ def _pre_load_state_dict_hook(
                 continue
 
             key = f"{prefix}embedding_bags.{table_name}.weight"
+
+            # If key not in state dict, continue
+            if key not in state_dict:
+                continue
+
             # gather model shards from both DTensor and ShardedTensor maps
             model_shards_sharded_tensor = self._model_parallel_name_to_local_shards[
                 table_name