24 lines
684 B
Python
24 lines
684 B
Python
import numpy as np
|
|
import xgboost as xgb
|
|
import cupy as cp
|
|
import time
|
|
import pytest
|
|
|
|
|
|
# Test for integer overflow or out of memory exceptions
|
|
def test_large_input():
|
|
available_bytes, _ = cp.cuda.runtime.memGetInfo()
|
|
# 15 GB
|
|
required_bytes = 1.5e+10
|
|
if available_bytes < required_bytes:
|
|
pytest.skip("Not enough memory on this device")
|
|
n = 1000
|
|
m = ((1 << 31) + n - 1) // n
|
|
assert (np.log2(m * n) > 31)
|
|
X = cp.ones((m, n), dtype=np.float32)
|
|
y = cp.ones(m)
|
|
dmat = xgb.DeviceQuantileDMatrix(X, y)
|
|
booster = xgb.train({"tree_method": "gpu_hist", "max_depth": 1}, dmat, 1)
|
|
del y
|
|
booster.inplace_predict(X)
|