pytorch_xgboost.py

import cupy
import torch
import GPUtil
import xgboost as xgb

from torch.utils.dlpack import to_dlpack
from torch.utils.dlpack import from_dlpack

mem_before = GPUtil.getGPUs()[0].memoryUsed
# Create a PyTorch tensor.
tx1 = torch.randn(1000000, 1000).cuda()

print(GPUtil.getGPUs()[0].memoryUsed - mem_before)
# Convert it into a DLPack tensor.
dx = to_dlpack(tx1)
# Train with xgboost!
dmat = xgb.DeviceQuantileDMatrix(dx)
print(GPUtil.getGPUs()[0].memoryUsed - mem_before)