the set of expected input keys.
    """
    # When running in graph mode (TF1), tf.Tensor types should be converted to
    # numpy ndarray types to be compatible with `make_tensor_proto`.
    if not context.executing_eagerly():
      with session.Session() as sess:
        repr_ds = replace_tensors_by_numpy_ndarrays(repr_ds, sess)

    expected_input_keys = self.expected_input_key_map.get(