Better message when no GPU is found. (#5594)

This commit is contained in:
Jiaming Yuan
2020-04-26 10:00:57 +08:00
committed by GitHub
parent 8dfe7b3686
commit 7d93932423
4 changed files with 37 additions and 0 deletions

View File

@@ -157,3 +157,14 @@ class TestPickling(unittest.TestCase):
bst.set_param({'predictor': 'cpu_predictor'})
cpu_pred = model.predict(x, output_margin=True)
np.testing.assert_allclose(cpu_pred, gpu_pred, rtol=1e-5)
def test_training_on_cpu_only_env(self):
cuda_environment = {'CUDA_VISIBLE_DEVICES': '-1'}
env = os.environ.copy()
env.update(cuda_environment)
args = self.args_template.copy()
args.append(
"./tests/python-gpu/"
"load_pickle.py::TestLoadPickle::test_training_on_cpu_only_env")
status = subprocess.call(args, env=env)
assert status == 0