fix broken test

Summary: the assertion causes test failure due to Proxy(KJT) instead of KJT Differential Revision: D66321923
TroyGarden · Nov 21, 2024 · e365130 · e365130
1 parent 54362a0
commit e365130
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/torchrec/distributed/test_utils/test_model.py b/torchrec/distributed/test_utils/test_model.py
@@ -1803,11 +1803,10 @@ def forward(
             modified_input = self._preproc_module(modified_input)
         elif self._run_preproc_inline:
             idlist_features = modified_input.idlist_features
-            assert isinstance(idlist_features, KeyedJaggedTensor)
             modified_input.idlist_features = KeyedJaggedTensor.from_lengths_sync(
-                idlist_features.keys(),
-                idlist_features.values(),
-                idlist_features.lengths(),
+                idlist_features.keys(),  # pyre-ignore [6]
+                idlist_features.values(),  # pyre-ignore [6]
+                idlist_features.lengths(),  # pyre-ignore [16]
             )
 
         modified_idlist_features = self.preproc_nonweighted(