remove hardcoded n_epochs

[johnlyzhou]

remove the hardcoded conditional requiring n_epochs==200

[themattinthehatt]

@johnlyzhou good catch here. We'll still need to add some additional updates though - filtering by n_epochs ensures that we only plot the loss values at the end of training; not doing so will plot the loss values for all epochs, which we don't want. So you'll have to find the max number of epochs in the dataframe, then filter by that value.