Added time step code to LSTM