Use strategy.scope()

official/resnet/keras/keras_cifar_main.py

@@ -132,16 +132,16 @@ def run(flags_obj):
       num_epochs=flags_obj.train_epochs,
       parse_record_fn=parse_record_keras)
 
-  optimizer = keras_common.get_optimizer()
   strategy = distribution_utils.get_distribution_strategy(
       flags_obj.num_gpus, flags_obj.turn_off_distribution_strategy)
 
-  model = resnet_cifar_model.resnet56(classes=cifar_main.NUM_CLASSES)
+  with strategy.scope():
+    optimizer = keras_common.get_optimizer()
+    model = resnet_cifar_model.resnet56(classes=cifar_main.NUM_CLASSES)
 
-  model.compile(loss='categorical_crossentropy',
-                optimizer=optimizer,
-                metrics=['categorical_accuracy'],
-                distribute=strategy)
+    model.compile(loss='categorical_crossentropy',
+                  optimizer=optimizer,
+                  metrics=['categorical_accuracy'])
 
   time_callback, tensorboard_callback, lr_callback = keras_common.get_callbacks(
       learning_rate_schedule, cifar_main.NUM_IMAGES['train'])

official/resnet/keras/keras_imagenet_main.py

@@ -121,16 +121,15 @@ def run(flags_obj):
                                 num_epochs=flags_obj.train_epochs,
                                 parse_record_fn=parse_record_keras)
 
-  optimizer = keras_common.get_optimizer()
   strategy = distribution_utils.get_distribution_strategy(
       flags_obj.num_gpus, flags_obj.turn_off_distribution_strategy)
 
-  model = resnet_model.resnet50(num_classes=imagenet_main.NUM_CLASSES)
-
-  model.compile(loss='sparse_categorical_crossentropy',
-                optimizer=optimizer,
-                metrics=['sparse_categorical_accuracy'],
-                distribute=strategy)
+  with strategy.scope():
+    optimizer = keras_common.get_optimizer()
+    model = resnet_model.resnet50(num_classes=imagenet_main.NUM_CLASSES)
+    model.compile(loss='sparse_categorical_crossentropy',
+                  optimizer=optimizer,
+                  metrics=['sparse_categorical_accuracy'])
 
   time_callback, tensorboard_callback, lr_callback = keras_common.get_callbacks(
       learning_rate_schedule, imagenet_main.NUM_IMAGES['train'])