import tensorflow as tf

        infer_step, [result, logits, loss],
        shape_invariants=[
            tf.TensorShape([None, None, None, None]),
            tf.TensorShape([None, None, None, None, None]),
