Switch division to tf.float32

yashk2810 · copybara-github · commit 8ca44ddefdbd · 2019-06-18T08:04:21.000-07:00
PiperOrigin-RevId: 253790284
diff --git a/tensorflow_examples/models/densenet/distributed_train.py b/tensorflow_examples/models/densenet/distributed_train.py
@@ -122,8 +122,8 @@ def custom_loop(self, train_dist_dataset, test_dist_dataset,
     """
 
     def distributed_train_epoch(ds):
-      total_loss = 0.
-      num_train_batches = 0
+      total_loss = 0.0
+      num_train_batches = 0.0
       for one_batch in ds:
         per_replica_loss = strategy.experimental_run_v2(
             self.train_step, args=(one_batch,))
@@ -133,7 +133,7 @@ def distributed_train_epoch(ds):
       return total_loss, num_train_batches
 
     def distributed_test_epoch(ds):
-      num_test_batches = 0
+      num_test_batches = 0.0
       for one_batch in ds:
         strategy.experimental_run_v2(
             self.test_step, args=(one_batch,))
diff --git a/tensorflow_examples/models/nmt_with_attention/distributed_train.py b/tensorflow_examples/models/nmt_with_attention/distributed_train.py
@@ -67,8 +67,8 @@ def training_loop(self, train_dist_dataset, test_dist_dataset,
     """
 
     def distributed_train_epoch(ds):
-      total_loss = 0.
-      num_train_batches = 0
+      total_loss = 0.0
+      num_train_batches = 0.0
       for one_batch in ds:
         per_replica_loss = strategy.experimental_run_v2(
             self.train_step, args=(one_batch,))
