Switching to a simpler test for keras examples for MWMS

Author: Lokiiiiii

tests/data/tensorflow_mnist/mnist_mwms.py

@@ -53,3 +53,5 @@ def build_and_compile_cnn_model():
     multi_worker_model = build_and_compile_cnn_model()
 
 multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70)
+
+print(f"strategy.num_replicas_in_sync={strategy.num_replicas_in_sync}")

tests/integ/test_tf.py

@@ -212,7 +212,7 @@ def test_mwms_gpu(
     logs = captured.out + captured.err
     print(logs)
     assert "Running distributed training job with multi_worker_mirrored_strategy setup" in logs
-    assert f"nranks {instance_count}" in logs
+    assert f"strategy.num_replicas_in_sync={instance_count}" in logs
 
 
 @pytest.mark.release