wandb init earlier

[aaprasad]

Currently, we only run wandb.init right before starting the actual train loop. It might be nice to init earlier so that the other outputs of the train script are logged to wandb especially in the case of running job arrays on the cluster where sometimes stdout can get lost