How to change train_batch to return additional metrics

We will log any key/value pairs returned from train_batch - need to show users that this is the case. 

Additionally, I believe there is an epoch level metric return too. IIRC this was a special quick fix made for applecider folks. Double check and document this in the same notebook. 