tf.data.Dataset.group_by_window() в TensorSpec

Question

я пытаюсь бежать tf.data.Dataset.group_by_window()в наборе данных, который у меня есть из petastorm make_tf_datasetи продолжайте получать ошибку

      ValueError: Invalid `key_func`. `key_func` must return a single `tf.int64` scalar tensor but its return type is TensorSpec(shape=(None,), dtype=tf.int64, name=None).

Код:

      with test_converter.make_tf_dataset(batch_size=BATCH_SIZE, num_epochs=1) as test_dataset:

  tf_test = test_dataset.map(row_generator, num_parallel_calls=tf.data.AUTOTUNE, deterministic=False)  

  key_func = lambda x: x["my_id_int"]
  reduce_func = lambda key, dataset: dataset.batch(100)
  tf_test_grp = tf_test.group_by_window(
      key_func=key_func, reduce_func=reduce_func, window_size=100)

генератор строк

      def row_generator(x):
  d = {'my_id_int':x.my_id_int, ...}
  return d

и test_converter

      test_converter = make_spark_converter(df_test.select(all_fields))

Кто-нибудь знает, как обойти это? Должен ли ключ или карта иметь разные значения?