Failed to call ThenRnnForward with model config

When I try to train my model, I get the following error:

tensorflow.python.framework.errors_impl.InternalError:    Failed to call ThenRnnForward with model config: [rnn_mode, rnn_input_mode, rnn_direction_mode]: 2, 0, 0 , [num_layers, input_size, num_units, dir_count, max_seq_length, batch_size, cell_num_units]: [1, 32, 32, 1, 4096, 256, 32]
         [[]]
         [[functional_1/lstm_2/PartitionedCall]] [Op:__inference_train_function_8148]

Function call stack:
train_function -> train_function -> train_function

I'm trying to train a simple LSTM model using the Tensorflow functional API.

Here is my model:

inp = Input(shape = (size,1))

x = LSTM(32, return_sequences=True)(inp)
x = LSTM(32, return_sequences=True)(x)
x = LSTM(32)(x)

x = Dense(512, kernel_initializer='he_uniform', use_bias = True)(x)
x = Activation('relu')(x)

x = Dense(256, kernel_initializer='he_uniform', use_bias = True)(x)
x = Activation('relu')(x)

x = Dense(12, activation = 'sigmoid')(x)

Where "size" is 4096. My data has the shape (num_samples, 4096, 1).

EDIT: Reducing the batch size significantly (from 512 to 64) works, but the training is incredibly slow, even on a GPU.

from Failed to call ThenRnnForward with model config