Add wandb support.

[ebezzam]

Thank you for submitting a pull request (PR)! Before hitting the submit button, we would be really glad if you check all the following points. Even if you can't check all the boxes below, do not hesitate to go ahead with the PR and ask for help there.

Please also refer to the doc on contributing for more details.

• Have you run the tests by running pytest from the repository root? See the documentation for more details.
• Do you follow the code formatting for this project? This can be done by setting up pre-commit hooks as described in the documentation.
• Is there a unit test for the proposed code modification? If the PR addresses an issue, the test should make sure the issue is fixed.
• Are there docstrings? Do they follow the numpydoc style?
• Have you checked that the doc builds properly and that any new RST file has been added to the repo? How to do that is covered in the documentation.
• Did you document the proposed change(s) in the CHANGELOG file? It should go under "Unreleased".

Happy PR 😃