From ca0ec779bf63a12147343c7b9eccb8f441cede67 Mon Sep 17 00:00:00 2001 From: Rahul Huilgol Date: Fri, 22 Nov 2019 14:47:28 -0800 Subject: [PATCH] Fix tests --- tests/tensorflow/keras/test_keras.py | 8 ++++---- tests/tensorflow/keras/test_keras_mirrored.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/tensorflow/keras/test_keras.py b/tests/tensorflow/keras/test_keras.py index b46e3789a..a100feb0b 100644 --- a/tests/tensorflow/keras/test_keras.py +++ b/tests/tensorflow/keras/test_keras.py @@ -84,10 +84,10 @@ def train_model( reduction_config=reduction_config, ) - if not save_all and include_collections is not None: - for cname in hook.include_collections: - if cname not in include_collections: - hook.get_collection(cname).save_config = SaveConfig(end_step=0) + if not save_all and include_collections is not None: + for cname in hook.include_collections: + if cname not in include_collections: + hook.get_collection(cname).save_config = SaveConfig(end_step=0) if create_relu_collection: hook.get_collection("relu").add_keras_layer(relu_layer, inputs=True, outputs=True) diff --git a/tests/tensorflow/keras/test_keras_mirrored.py b/tests/tensorflow/keras/test_keras_mirrored.py index 94991a444..71b5e6923 100644 --- a/tests/tensorflow/keras/test_keras_mirrored.py +++ b/tests/tensorflow/keras/test_keras_mirrored.py @@ -123,10 +123,10 @@ def scale(image, label): include_workers=include_workers, ) - if not save_all and include_collections is not None: - for cname in hook.include_collections: - if cname not in include_collections: - hook.get_collection(cname).save_config = SaveConfig(end_step=0) + if not save_all and include_collections is not None: + for cname in hook.include_collections: + if cname not in include_collections: + hook.get_collection(cname).save_config = SaveConfig(end_step=0) if use_keras_optimizer: opt = tf.keras.optimizers.Adam()