Обучение модели seq2seq в Google Colab TPU с большим набором данных - Керас

Я пытаюсь обучить последовательность для модели последовательности для машинного перевода с помощью Keras в Google Colab TPU. У меня есть набор данных, который я могу загрузить в память, но мне нужно предварительно обработать его, чтобы передать его модели. В частности, мне нужно преобразовать целевые слова в одни горячие векторы, и со многими примерами я не могу загрузить все преобразование в память, поэтому мне нужно создавать пакеты данных.

Я использую эту функцию как пакетный генератор:

def generate_batch_bert(X_ids, X_masks, y, batch_size = 1024):
    ''' Generate a batch of data '''
    while True:
        for j in range(0, len(X_ids), batch_size):
          # batch of encoder and decoder data
          encoder_input_data_ids = X_ids[j:j+batch_size]
          encoder_input_data_masks = X_masks[j:j+batch_size]
          y_decoder = y[j:j+batch_size]
          

          # decoder target and input for teacher forcing
          decoder_input_data = y_decoder[:,:-1]
          decoder_target_seq = y_decoder[:,1:]
          
          # batch of decoder target data
          decoder_target_data = to_categorical(decoder_target_seq, vocab_size_fr)
          # keep only with the right amount of instances for training on TPU
          if encoder_input_data_ids.shape[0] == batch_size:
            yield([encoder_input_data_ids, encoder_input_data_masks, decoder_input_data], decoder_target_data)

Проблема в том, что всякий раз, когда я пытаюсь запустить функцию соответствия, как показано ниже:

model.fit(x=generate_batch_bert(X_train_ids, X_train_masks, y_train, batch_size = batch_size),
                    steps_per_epoch = train_samples//batch_size,
                    epochs=epochs,
                    callbacks = callbacks,
                    validation_data = generate_batch_bert(X_val_ids, X_val_masks, y_val, batch_size = batch_size),
                    validation_steps = val_samples//batch_size)

Я получаю следующую ошибку:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/tensor_util.py:445 make_tensor_proto
    raise ValueError("None values not supported.")

ValueError: None values not supported.

Не уверен, что случилось и как я могу решить эту проблему.

ИЗМЕНИТЬ

Я попытался загрузить в память меньшее количество данных, чтобы преобразование в одну горячую кодировку целевых слов не приводило к сбою ядра и оно действительно работало. Итак, очевидно, что что-то не так с тем, как я генерирую партии.