I could reproduce this with a simple example. At the first glance it looks like a bug in xgboost-ray’s prediction code. I’ll take a look.
As a workaround and if feasible (if the data is small enough), can you pass distributed=False to your RayDMatrix creation?
pred_set = RayDMatrix(data = ray_dataset,
label = "rz_flag",
distributed=False
)