diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 5817b484aa..ef4c2fd0e9 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -538,6 +538,8 @@ def train_retinanet(): print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m") print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, " f"epoch global_mem: {steps_in_train_epoch * GlobalCounters.global_mem:_}") + # if we are doing beam search, run the first eval too + if (TRAIN_BEAM or EVAL_BEAM) and e == start_epoch: break return # ** eval loop **