support TPU on resnet50 using keras compile/fit

PiperOrigin-RevId: 273358759
hz7715 · Oct 7, 2019 · fa5b66e · fa5b66e
1 parent 482784a
commit fa5b66e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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