test low batch size

examples/jax/jax-dist-spmd-mnist/spmd_mnist_classifier_fromscratch.py

@@ -99,7 +99,8 @@ def accuracy(params, batch):
     # For this manual SPMD example, we get the number of devices (e.g. CPU,
     # GPUs or TPU cores) that we're using, and use it to reshape data minibatches.
     num_devices = jax.local_device_count()
-    batch_size = num_devices * 5
+    # batch_size = num_devices * 5
+    batch_size = 5 # testing
 
     train_images, train_labels, test_images, test_labels = datasets.mnist()
     num_train = train_images.shape[0]