Apply LSTM to each matrix element with Keras
SOLUTION

Replace:

# Apply if statement
y = keras.backend.switch(cond, 
                          lambda: keras.layers.GRU(units,return_sequences=True)(seq), 
                          lambda: keras.backend.zeros(shape=(1,time_steps,units)))

y = keras.layers.Lambda(convert_to_layer)(y)


with:

y = keras.layers.Lambda(lambda x: keras.backend.switch(cond, 
                          lambda: keras.layers.GRU(units,return_sequences=True)(x), 
                          lambda: keras.backend.zeros(shape=(1,time_steps,units))))(seq)