Проблема с эпохой keras fit_generator

Question

Я создаю модель LSTM для генерации текста с использованием Keras. Поскольку набор данных (около 25 романов, содержащий около 1,4 миллиона слов), который я использую, не может быть обработан сразу (проблема с памятью при преобразовании моих выходных данных в _Categorical ()), я создал собственную функцию генератора для чтения данных.

# Data generator for fit and evaluate
def generator(batch_size):
    start = 0
    end = batch_size
    while True:
      x = sequences[start:end,:-1]
      #print(x)
      y = sequences[start:end,-1]
      y = to_categorical(y, num_classes=vocab_size)
      #print(y)
      yield x, y
      if batch_size == len(lines):
        break;
      else:
        start += batch_size
        end += batch_size

когда я выполняю метод model.fit (), после обучения 1 эпохи выдается следующая ошибка.

UnknownError:  [_Derived_]  CUDNN_STATUS_BAD_PARAM
in tensorflow/stream_executor/cuda/cuda_dnn.cc(1459): 'cudnnSetTensorNdDescriptor( tensor_desc.get(), data_type, sizeof(dims) / sizeof(dims[0]), dims, strides)'
     [[{{node CudnnRNN}}]]
     [[sequential/lstm/StatefulPartitionedCall]] [Op:__inference_train_function_25138]

Function call stack:
train_function -> train_function -> train_function

кто-нибудь знает, как решить эту проблему? Спасибо