This code keeps on running till we close the terminal Can we set the maximum epoch till which the training should run?