Adding WandB support

[prajdabre]

Currently, the model information such as losses and gradients are saved every N steps to a local file, which is then visualized using tensorboard. WandB is the cool new kid on the block, but we must respect our elders.

Adding a flag and then placing if else to switch between wandb and tensorboard would be nice. If someone wants to use wandb then remember to print a message so that users specify their username and wandb workspace url via flags. Also, in case of a wandb error, asking users to check if the wandb initialization is done properly or not.

This is fairly easy. As usual, please document.