wandb and xgboost

import wandb
import numpy as np
import xgboost as xgb
from keras.datasets import mnist 

def wandb_callback():
    def callback(env):
        for k, v in env.evaluation_result_list:
            wandb.log({k: v}, commit=False)
        wandb.log({})

    return callback


# logging code
run = wandb.init()
config = run.config

(X_train, y_train), (X_test, y_test) = mnist.load_data()
X_train = X_train.reshape(X_train.shape[0], -1)
X_test = X_test.reshape(X_test.shape[0], -1)

param_list = [("eta", 0.08), ("max_depth", 6), ("subsample", 0.8), ("colsample_bytree", 0.8), ("objective", "multi:softmax"), ("eval_metric", "merror"), ("alpha", 8), ("lambda", 2), ("num_class", 10)]
config.update(dict(param_list))

config.n_rounds = 2
config.early_stopping = 50

d_train = xgb.DMatrix(X_train, label=y_train)
d_val = xgb.DMatrix(X_test, label=y_test)
eval_list = [(d_train, "train"), (d_val, "validation")]
bst = xgb.train(param_list, d_train, config.n_rounds, evals=eval_list, early_stopping_rounds=config.early_stopping, verbose_eval=True, callbacks=[wandb_callback()])