Fix generate_vbe_metadata in DenseTBE (#3095)

Summary:
X-link: facebookresearch/FBGEMM#187

Pull Request resolved: #3095

As title

Reviewed By: q10, spcyppt

Differential Revision: D62387892

fbshipit-source-id: cfc85587891b1d7e87920768540218d31fb5775f

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -2922,17 +2922,6 @@ def _generate_vbe_metadata(
     ) -> invokers.lookup_args.VBEMetadata:
         # Blocking D2H copy, but only runs at first call
         self.feature_dims = self.feature_dims.cpu()
-        if batch_size_per_feature_per_rank is not None:
-            assert self.optimizer in (
-                OptimType.EXACT_ROWWISE_ADAGRAD,
-                OptimType.EXACT_SGD,
-                OptimType.ENSEMBLE_ROWWISE_ADAGRAD,
-                OptimType.NONE,
-            ), (
-                "Variable batch size TBE support is enabled for "
-                "OptimType.EXACT_ROWWISE_ADAGRAD and "
-                "ENSEMBLE_ROWWISE_ADAGRAD only"
-            )
         return generate_vbe_metadata(
             offsets,
             batch_size_per_feature_per_rank,