Fix wrong return type of QuantManagedCollisionEmbeddingCollection (#2830)

Summary:
Pull Request resolved: #2830

this forward will be called after diff D68991644, so we need to make sure the return type is consistent with ManagedCollisionEmbeddingCollection

Reviewed By: kausv, jma99fb

Differential Revision: D71250093

fbshipit-source-id: a5c91a2f703fe1803ebcc54be9cb3634bcec1354

Author: emlin

torchrec/distributed/tests/test_infer_shardings.py

@@ -2229,7 +2229,7 @@ def test_sharded_quant_mc_ec_rw(
         sharded_model.load_state_dict(quant_model.state_dict())
         sharded_output = sharded_model(*inputs[0])
 
-        assert_close(non_sharded_output, sharded_output)
+        assert_close(non_sharded_output[0], sharded_output[0])
         gm: torch.fx.GraphModule = symbolic_trace(
             sharded_model,
             leaf_modules=[
@@ -2242,7 +2242,7 @@ def test_sharded_quant_mc_ec_rw(
         gm_script = torch.jit.script(gm)
         print(f"gm_script:\n{gm_script}")
         gm_script_output = gm_script(*inputs[0])
-        assert_close(sharded_output, gm_script_output)
+        assert_close(sharded_output[0], gm_script_output[0])
 
     @unittest.skipIf(
         torch.cuda.device_count() <= 1,

torchrec/quant/embedding_modules.py

@@ -1093,7 +1093,9 @@ def forward(
     ) -> Dict[str, JaggedTensor]:
         features = self._managed_collision_collection(features)
 
-        return super().forward(features)
+        # mcec expects Tuple return type
+        # pyre-ignore
+        return (super().forward(features),)
 
     def _get_name(self) -> str:
         return "QuantManagedCollisionEmbeddingCollection"