Look into adding logging

**Describe the solution you'd like**
Logging would be quite useful, to diagnose problems with the library itself. Logging could also represent TensorBoard/wandb and other experiment trackers.

**Additional context**
It would be helpful to hook `SummaryWriter` or some other way to log experiments automatically.
