NameError when wandb is not found

@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.

Start by comparing __try_wandb in lanfactory/trainers/jax_mlp.py and torch_mlp.py, focusing on the missing wandb name shown in the traceback. Reproduce the JAX trainer with wandb unavailable and verify that it no longer raises this NameError while the existing torch trainer behavior remains unchanged.

Written by the indexing model from the issue text.