NameError when wandb is not found

[Jasonleng]

@AlexanderFengler When I try the jax trainer it get the error message:

Traceback (most recent call last):
File "/oscar/home/xleng/LAN_pipeline_minimal/bash_scripts/../scripts/jax_training_script.py", line 185, in
model_trainer.train_and_evaluate(save_history = train_config['save_history'],
File "/users/xleng/anaconda/lan_pipe/lib/python3.10/site-packages/lanfactory/trainers/jax_mlp.py", line 542, in train_and_evaluate
self.__try_wandb(
File "/users/xleng/anaconda/lan_pipe/lib/python3.10/site-packages/lanfactory/trainers/jax_mlp.py", line 365, in __try_wandb
wandb.init(
NameError: name 'wandb' is not defined

It does not happen when I do the torch trainer. And I realized that the implementation of the __try_wandb function varies between jax_mlp.py and torch_mlp.py.