Follow the original example in specifying input/output shape

Author: abheesht17

examples/ml_perf/configs/v6e_8.py

@@ -196,6 +196,7 @@
 # Set `num_steps` in the main config file instead of num_epochs, because we are
 # using a Python generator.
 training_config.num_steps = 2
+training_config.eval_freq = 1
 
 # === Assign all configs to the root config ===
 config.dataset = dataset_config

examples/ml_perf/main.py

@@ -84,9 +84,9 @@ def main(
             name=feature_name,
             table=table_config,
             # TODO: Verify whether it should be `(bsz, 1)` or
-            # `(bsz, feature_list_length)`.
-            input_shape=(per_host_batch_size, feature_list_length),
-            output_shape=(per_host_batch_size, embedding_dim),
+            # `(bsz, feature_list_length)`. The original example uses 1.
+            input_shape=(global_batch_size, 1),
+            output_shape=(global_batch_size, embedding_dim),
         )
 
     # === Instantiate model ===