Thanks for the reply. It slows down during the second epoch. There are about 400 time steps to compute for each epoch; in epoch 1, it takes about equal amount of time to compute each step, but in epoch 2 the time it takes to compute each step starts to increase at around step 300. And in fact I’ve run the model overnight but the optimizer.step() call in epoch 2 hasn’t returned yet. So, I can’t tell whether epoch 3 will take longer because I haven’t reached there yet.