[FEATURE]: wandb model checkpoint support

**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 
