Training an RNN model in TensorFlow