Initial work on tracing callback

[dfalbel]

THis callback will trace the model for training and validation. Hopefully the optimizations that happen on tracing can generate some speedups for training more complicated models.

TODO:

• [ ] correctly share weights between training and validation