remove the unnecessary fx wrapped function in test_model (#2900)

Summary:
Pull Request resolved: #2900

# context
* prevously this `_post_ebc_test_wrap_function` function is added because of fx tracing issue D45133173
> Currently when fx tracing ebc would output `EmbeddingBagCollectionAwaitable` instead of a proxy object. If there're a wrapped function which takes this as input (and neither of the other inputs is proxy), then we would see such error

* This problem is gone so this wrap function is no longer needed

Reviewed By: dstaay-fb

Differential Revision: D63901161

fbshipit-source-id: c2e6694754a31d100e24a7da573aa58f3befcbd4

Author: TroyGarden

torchrec/distributed/test_utils/test_model.py

@@ -1285,7 +1285,6 @@ def forward(
             for fp in self.fps:
                 fp_features = fp(fp_features)
         ebc = self.ebc(features)
-        ebc = _post_ebc_test_wrap_function(ebc)
         fp_ebc: Optional[KeyedTensor] = (
             self.fp_ebc(fp_features) if self.fp_ebc is not None else None
         )
@@ -1911,14 +1910,6 @@ def _get_default_rtol_and_atol(
     return max(actual_rtol, expected_rtol), max(actual_atol, expected_atol)
 
 
-@torch.fx.wrap
-def _post_ebc_test_wrap_function(kt: KeyedTensor) -> KeyedTensor:
-    for _ in kt.values():
-        continue
-
-    return kt
-
-
 class TestPreprocNonWeighted(nn.Module):
     """
     Basic module for testing
@@ -2374,7 +2365,6 @@ def forward(
             KeyedTensor
         """
         ebc, _ = self.ebc(features)
-        ebc = _post_ebc_test_wrap_function(ebc)
         w_ebc, _ = (
             self.weighted_ebc(weighted_features)
             if self.weighted_ebc is not None and weighted_features is not None