Using Wandb with tune.run()

I am having issues training my environment and visualizing with wandb. I am following this code but it’s not working for me?

from ray.tune.logger import DEFAULT_LOGGERS
from ray.tune.integration.wandb import WandbLogger
tune.run(
    train_fn,
    config={
        # define search space here
        "parameter_1": tune.choice([1, 2, 3]),
        "parameter_2": tune.choice([4, 5, 6]),
        # wandb configuration
        "wandb": {
            "project": "Optimization_Project",
            "api_key_file": "/path/to/file",
            "log_config": True
        }
    },
    loggers=DEFAULT_LOGGERS + (WandbLogger, ))

Nevermind, resolved! It worked after I switched to WandbLoggerCallback:

results = tune.run(
    alg, 
    stop={"training_iteration": 10}, 
    config=config, 
    local_dir="results/test_viz/{}".format(date),
    callbacks=[
        WandbLoggerCallback(
            api_key="123456",  # Censored for obvious reasons
            project="cartpole_viz"
        )
    ]
)