Introduction

When a TensorFlow custom Estimator appears to hang during evaluate() right after training, the culprit is usually not the evaluator itself. In practice, the problem is often a non-terminating eval dataset, a model_fn that mishandles eval mode, or checkpoint confusion between the training and evaluation steps. The fastest fix is to debug evaluation as a separate pipeline, not as an extension of training.

Make Sure the Eval Input Is Finite

The most common cause is an eval dataset that never ends. Estimator evaluation stops when either the dataset is exhausted or the requested steps count is reached. If the eval input repeats forever and you do not pass steps, evaluate() has no reason to return.

1import tensorflow as tf
2
3def eval_input_fn():
4    features = tf.constant([[1.0], [2.0], [3.0], [4.0]])
5    labels = tf.constant([0, 0, 1, 1])
6
7    dataset = tf.data.Dataset.from_tensor_slices(({"x": features}, labels))
8    dataset = dataset.batch(2)
9    return dataset

What you do not want is an eval path like this:

dataset = dataset.repeat()

That line is fine for training, but it is a common reason evaluation appears stuck.

Keep the model_fn Branches Clean

A custom estimator uses model_fn(features, labels, mode, params) for train, eval, and predict. Each mode should return only the pieces relevant to that mode.

1import tensorflow as tf
2
3def model_fn(features, labels, mode, params):
4    x = features["x"]
5    logits = tf.keras.layers.Dense(2)(x)
6    predictions = tf.argmax(logits, axis=1, output_type=tf.int32)
7
8    if mode == tf.estimator.ModeKeys.PREDICT:
9        return tf.estimator.EstimatorSpec(
10            mode=mode,
11            predictions={"class_ids": predictions},
12        )
13
14    loss = tf.reduce_mean(
15        tf.nn.sparse_softmax_cross_entropy_with_logits(
16            labels=labels,
17            logits=logits,
18        )
19    )
20
21    metrics = {
22        "accuracy": tf.compat.v1.metrics.accuracy(
23            labels=labels,
24            predictions=predictions,
25        )
26    }
27
28    if mode == tf.estimator.ModeKeys.EVAL:
29        return tf.estimator.EstimatorSpec(
30            mode=mode,
31            loss=loss,
32            eval_metric_ops=metrics,
33        )
34
35    optimizer = tf.compat.v1.train.AdamOptimizer(0.01)
36    train_op = optimizer.minimize(loss, tf.compat.v1.train.get_global_step())
37    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

If eval mode accidentally depends on train hooks, infinite input, or side effects intended only for training, the hang can be hard to diagnose from the outside.

Debug with Explicit Eval Steps

During troubleshooting, force evaluation to run for a small fixed number of steps.

1results = estimator.evaluate(
2    input_fn=eval_input_fn,
3    steps=2,
4)
5print(results)

If that works but steps=None does not, the dataset is likely the issue. If even a small fixed step count hangs, the bug is probably inside the input function or model_fn.

Verify Checkpoints and Dataset Independently

Do not assume the data path and checkpoint path are both correct just because training succeeded. Confirm each one separately.

1import tensorflow as tf
2
3print(tf.train.latest_checkpoint("/tmp/my_estimator_model"))
4
5for batch in eval_input_fn().take(2):
6    print(batch)

Those two checks answer different questions:

• is there a valid checkpoint to load
• can the eval dataset actually produce batches

If the dataset blocks when iterated directly, the estimator is not the real problem.

Watch Out for Input Functions That Hide Blocking Work

Evaluation hangs are especially common when the input pipeline uses:

• 'Dataset.from_generator'
• 'tf.py_function'
• slow remote filesystems
• background preprocessing threads

Those features can block quietly if the generator never ends, an external process stalls, or preprocessing deadlocks. Strip the pipeline down to constant tensors first, then add complexity back one piece at a time.

Common Pitfalls

The biggest mistake is reusing the training dataset for evaluation without removing .repeat(). That single line creates an infinite eval stream and makes evaluate() appear frozen.

Another issue is writing a model_fn that mixes mode-specific behavior. Train-only hooks, optimizer dependencies, or logging side effects can leak into eval mode if the branches are not cleanly separated.

Checkpoint confusion is also common. If model_dir points somewhere unexpected, evaluation may wait on the wrong files or reload stale state, which makes the failure mode look unrelated to checkpoints at first glance.

Finally, do not debug only through the Estimator API. Test the eval dataset as a standalone iterator and verify the latest checkpoint directly. That narrows the search space much faster than treating the whole training loop as one black box.

Summary

• A stuck evaluate() call is usually caused by the input pipeline or mode-specific logic, not by the evaluator alone.
• Make sure the eval dataset is finite unless you pass an explicit steps value.
• Keep train and eval behavior separate inside model_fn.
• Verify checkpoints and iterate the eval dataset independently during debugging.
• Reduce the pipeline to a minimal reproducible example before adding real data and complexity back.