Add model graph to logger

[jerpint]

add graph to visualize models in the logger (tensorboard in this case)

[jerpint]

It might be a good idea to integrate aim support before merging this in