signature_key,
          tags,
          inputs={input_key: in_placeholder},
          outputs={output_key: output_tensor},
      )

    return in_placeholder

  def _create_while_model(self, input_shape: Sequence[int] = (1, 32, 32, 512)):
    class WhileModel(module.Module):
      """A model with a while op."""

      def __init__(self):