removed a consider-using-in pitfall case

EleutherAI · Oct 24, 2021 · 8f55f6c · 8f55f6c
1 parent 7fbd31b
commit 8f55f6c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model_fns.py b/model_fns.py
@@ -135,7 +135,7 @@ def scaffold_fn():
             prediction_hooks=[mtf.MtfRestoreHook(lowering)])
 
     # We're not predicting, so we better be training or evaluating
-    assert (mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL)
+    assert mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL]
 
     if mode == tf.estimator.ModeKeys.TRAIN:
         # Gets number of microbatches per batch for serialized training