support TPU on resnet50 using keras compile/fit

PiperOrigin-RevId: 273358759

support TPU on resnet50 using keras compile/fit
PiperOrigin-RevId: 273358759
fa5b66e5 · Jing Li · A. Unique TensorFlower · 482784ab · fa5b66e5
Commit fa5b66e5 authored Oct 07, 2019 by Jing Li Committed by A. Unique TensorFlower Oct 07, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/vision/image_classification/resnet_imagenet_main.py official/vision/image_classification/resnet_imagenet_main.py +2 -1

No files found.
--- a/official/vision/image_classification/resnet_imagenet_main.py
+++ b/official/vision/image_classification/resnet_imagenet_main.py
@@ -83,7 +83,8 @@ def run(flags_obj):
      num_gpus=flags_obj.num_gpus,
      num_workers=num_workers,
      all_reduce_alg=flags_obj.all_reduce_alg,
-      num_packs=flags_obj.num_packs)
+      num_packs=flags_obj.num_packs,
+      tpu_address=flags_obj.tpu)

  if strategy:
    # flags_obj.enable_get_next_as_optional controls whether enabling