Bug fix: None check on strategy

official/resnet/keras/keras_cifar_main.py

@@ -135,7 +135,9 @@ def run(flags_obj):
   strategy = distribution_utils.get_distribution_strategy(
       flags_obj.num_gpus, flags_obj.turn_off_distribution_strategy)
 
-  with strategy.scope():
+  strategy_scope = keras_common.get_strategy_scope(strategy)
+
+  with strategy_scope:
     optimizer = keras_common.get_optimizer()
     model = resnet_cifar_model.resnet56(classes=cifar_main.NUM_CLASSES)
 

official/resnet/keras/keras_common.py

@@ -210,3 +210,22 @@ def get_synth_input_fn(height, width, num_channels, num_classes,
     return data
 
   return input_fn
+
+
+def get_strategy_scope(strategy):
+  if strategy:
+    strategy_scope = strategy.scope()
+  else:
+    strategy_scope = keras_common.DummyContextManager()
+
+  return strategy_scope
+
+
+class DummyContextManager(object):
+  def __enter__(self):
+    pass
+
+  def __exit__(self, *args):
+    pass
+
+

official/resnet/keras/keras_imagenet_main.py

@@ -124,9 +124,12 @@ def run(flags_obj):
   strategy = distribution_utils.get_distribution_strategy(
       flags_obj.num_gpus, flags_obj.turn_off_distribution_strategy)
 
-  with strategy.scope():
+  strategy_scope = keras_common.get_strategy_scope(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'])