Fix crash: "No module named 'tf_keras'".

Directly importing tensorflow.python.layers.normalization seems to no longer work.

PiperOrigin-RevId: 566806677

scripts/tf_cnn_benchmarks/convnet_builder.py

@@ -25,7 +25,6 @@
 import mlperf
 from tensorflow.python.layers import convolutional as conv_layers
 from tensorflow.python.layers import core as core_layers
-from tensorflow.python.layers import normalization as normalization_layers
 from tensorflow.python.layers import pooling as pooling_layers
 from tensorflow.python.training import moving_averages
 
@@ -462,7 +461,7 @@ def batch_norm(self, input_layer=None, decay=0.999, scale=False,
     center = True
     with tf.variable_scope(name) as scope:
       if self.use_tf_layers:
-        layer_obj = normalization_layers.BatchNormalization(
+        layer_obj = tf.layers.BatchNormalization(
             momentum=decay,
             scale=scale,
             epsilon=epsilon,