add test coverage for loading state dict from cpu for EBC sharding

torchrec/distributed/embeddingbag.py

@@ -597,10 +597,7 @@ def __init__(
         if env.process_group and dist.get_backend(env.process_group) != "fake":
             self._initialize_torch_state()
 
-        if module.device not in ["meta", "cpu"] and module.device.type not in [
-            "meta",
-            "cpu",
-        ]:
+        if module.device != "meta" and module.device.type != "meta":
             self.load_state_dict(module.state_dict(), strict=False)
 
     @classmethod

torchrec/distributed/planner/tests/test_types.py

@@ -419,8 +419,8 @@ def _test_hashing_consistency(
 
 
 class TestConsistentHashingBetweenProcesses(MultiProcessTestBase):
-
-    def test_hash_consistency(self) -> None:
+    # the proposal order might vary in github action so skip this test
+    def test_hash_consistency_disabled_in_oss_compatibility(self) -> None:
         # planner
         world_size = 2
         return_hash_dict = multiprocessing.Manager().dict()

torchrec/distributed/tests/test_init_parameters.py

@@ -223,7 +223,7 @@ def test_initialize_parameters_ec(
         device=st.sampled_from(
             [
                 torch.device("cuda"),
-                # torch.device("cpu"),
+                torch.device("cpu"),
             ]
         ),
     )