call zch forward with kjt weights to support feature score eviction (#3352)

Summary:
Pull Request resolved: #3352

as title

Reviewed By: steven1327, EddyLXJ

Differential Revision: D81726703

fbshipit-source-id: 9e1393da2f09404db7166c510b75a92b0d803a99

Author: emlin

torchrec/distributed/batched_embedding_kernel.py

@@ -2217,6 +2217,17 @@ def split_embedding_weights(
     ]:
         return self.emb_module.split_embedding_weights(no_snapshot, should_flush)
 
+    def forward(self, features: KeyedJaggedTensor) -> torch.Tensor:
+        # reset split weights during training
+        self._split_weights_res = None
+        self._optim.set_sharded_embedding_weight_ids(sharded_embedding_weight_ids=None)
+
+        return self.emb_module(
+            indices=features.values().long(),
+            offsets=features.offsets().long(),
+            weights=features.weights_or_none(),
+        )
+
 
 class ZeroCollisionEmbeddingCache(ZeroCollisionKeyValueEmbedding):
     def __init__(