diff --git a/tensorflow_privacy/privacy/keras_models/dp_keras_model.py b/tensorflow_privacy/privacy/keras_models/dp_keras_model.py
index b854f54..19b6fd9 100644
--- a/tensorflow_privacy/privacy/keras_models/dp_keras_model.py
+++ b/tensorflow_privacy/privacy/keras_models/dp_keras_model.py
@@ -1,3 +1,7 @@
+'''
+Keras Model for vectorized dpsgd with XLA acceleration
+'''
+
 import tensorflow as tf