def METHOD_NAME():
    ds = vaex.ml.iris()
    features = ['sepal_width', 'petal_length', 'sepal_length', 'petal_width']
    # Vanilla pygbm
    X = np.array(ds[features])
    # dtrain = lgb.Dataset(X, label=ds.data.class_)
    lgb_bst = lgb.train(param, dtrain, 3)
    lgb_pred = np.argmax(lgb_bst.predict(X), axis=1)
    # Vaex.ml.pygbm
    booster = ds.ml.pygbm_model(label='class_', num_round=3, features=features, param=param, classifier=True)
    vaex_pred = booster.predict(ds)
    # Comparing the the predictions of pygbm vs vaex.ml
    np.testing.assert_equal(vaex_pred, lgb_pred, verbose=True, err_msg='The predictions of vaex.ml do not match those of pygbm')