Addressed reviewer's comments.

Change-Id: I4e0c012f75813a18891072839da79c73c266ecad

Author: wwwind

tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_mha.py

@@ -39,7 +39,7 @@
 
 # define model
 input = tf.keras.layers.Input(shape=(28, 28))
-x = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=16, name="mha")(
+x = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=16, name='mha')(
     query=input, value=input
 )
 x = tf.keras.layers.Flatten()(x)
@@ -48,9 +48,9 @@
 
 # Train the digit classification model
 model.compile(
-    optimizer="adam",
+    optimizer='adam',
     loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-    metrics=["accuracy"],
+    metrics=['accuracy'],
 )
 
 model.fit(
@@ -80,9 +80,9 @@
 
 # `prune_low_magnitude` requires a recompile.
 model_for_pruning.compile(
-    optimizer="adam",
+    optimizer='adam',
     loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-    metrics=["accuracy"],
+    metrics=['accuracy'],
 )
 
 model_for_pruning.fit(
@@ -96,4 +96,4 @@
 
 score = model_for_pruning.evaluate(test_images, test_labels, verbose=0)
 print('Pruned model test loss:', score[0])
-print('Pruned model test accuracy:', score[1])
+print('Pruned model test accuracy:', score[1])