Fix total block num bug in tf2 efficientnet.

PiperOrigin-RevId: 314405729

Fix total block num bug in tf2 efficientnet.
PiperOrigin-RevId: 314405729
2af9696b · A. Unique TensorFlower · 7580f0a2 · 2af9696b
Commit 2af9696b authored Jun 02, 2020 by A. Unique TensorFlower
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/vision/image_classification/efficientnet/efficientnet_model.py ...n/image_classification/efficientnet/efficientnet_model.py +2 -1

No files found.
--- a/official/vision/image_classification/efficientnet/efficientnet_model.py
+++ b/official/vision/image_classification/efficientnet/efficientnet_model.py
@@ -372,7 +372,8 @@ def efficientnet(image_input: tf.keras.layers.Input,
                   name='stem')

  # Build blocks
-  num_blocks_total = sum(block.num_repeat for block in blocks)
+  num_blocks_total = sum(
+      round_repeats(block.num_repeat, depth_coefficient) for block in blocks)
  block_num = 0

  for stack_idx, block in enumerate(blocks):