Sampling callback

[rousseab]

A first crack at periodically sampling the model while training, through a callback.

I also refactor how the loss is logged, to make it super extra explicit what is logged at every step and what is aggregated at every epoch.

There's a lot of possible improvement as future work:

• there's not a lot of repetition in the input parameters... It would be nice to not have to repeat ourselves
• the histogram image that is logged is pretty bare boned: it would be nice to overlay over the validation data, but that's complicated to do right now.