From 6c705a052596e53b3adf758bdcbc2f3f243f9764 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Tue, 25 Jun 2019 18:10:15 -0400
Subject: [PATCH] adding framework level dp

---
 docs/source/examples/example_model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/source/examples/example_model.py b/docs/source/examples/example_model.py
index b946f487..e3917248 100644
--- a/docs/source/examples/example_model.py
+++ b/docs/source/examples/example_model.py
@@ -55,7 +55,7 @@ class ExampleModel(RootModule):
         nll = F.nll_loss(logits, labels)
         return nll
 
-    def training_step(self, data_batch):
+    def training_step(self, data_batch, batch_i):
         """
         Called inside the training loop
         :param data_batch:
@@ -72,7 +72,7 @@ class ExampleModel(RootModule):
         tqdm_dic = {'tng_loss': loss_val.item()}
         return loss_val, tqdm_dic
 
-    def validation_step(self, data_batch):
+    def validation_step(self, data_batch, batch_i):
         """
         Called inside the validation loop
         :param data_batch: