[FEATURE]: wandb model checkpoint support

[nmichlo]

Is your feature request related to a problem? Please describe. We don't just want to create checkpoints locally, but also upload to w&b

Describe the solution you'd like should extend the current check pointing system for wandb

Describe alternatives you've considered N/A

Additional context Extending #28