diff --git a/tensorflow_probability/python/bijectors/batch_normalization_test.py b/tensorflow_probability/python/bijectors/batch_normalization_test.py
index bb29345223..7a1604380d 100644
--- a/tensorflow_probability/python/bijectors/batch_normalization_test.py
+++ b/tensorflow_probability/python/bijectors/batch_normalization_test.py
@@ -220,7 +220,7 @@ def testInvertMutuallyConsistent(self, layer_cls):
 
   def testWithKeras(self):
     # NOTE: Keras throws an error below if we use
-    # tf1.layers.BatchNormalization() here.
+    # tf_keras.tf1_layers.BatchNormalization() here.
     layer = None
 
     dist = transformed_distribution.TransformedDistribution(