From a3e8a45559bade7841d4454cee267adf5894ea09 Mon Sep 17 00:00:00 2001
From: Shuang Song <shuangsong@google.com>
Date: Mon, 27 Feb 2023 11:11:27 -0800
Subject: [PATCH] Passes number of microbatches to DP model.

PiperOrigin-RevId: 512678620
---
 tutorials/mnist_dpsgd_tutorial_keras_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tutorials/mnist_dpsgd_tutorial_keras_model.py b/tutorials/mnist_dpsgd_tutorial_keras_model.py
index abd3683..b000c03 100644
--- a/tutorials/mnist_dpsgd_tutorial_keras_model.py
+++ b/tutorials/mnist_dpsgd_tutorial_keras_model.py
@@ -113,7 +113,9 @@ def main(unused_argv):
     model = DPSequential(
         l2_norm_clip=FLAGS.l2_norm_clip,
         noise_multiplier=FLAGS.noise_multiplier,
-        layers=layers)
+        num_microbatches=FLAGS.microbatches,
+        layers=layers,
+    )
   else:
     model = tf.keras.Sequential(layers=layers)