(self)
| 45 | assert config["learner"]["generic_param"]["device"] == "cuda:0" |
| 46 | |
| 47 | def test_wrap_gpu_id(self) -> None: |
| 48 | assert os.environ["CUDA_VISIBLE_DEVICES"] == "0" |
| 49 | bst = load_pickle(model_path) |
| 50 | config = bst.save_config() |
| 51 | config = json.loads(config) |
| 52 | assert config["learner"]["generic_param"]["device"] == "cuda:0" |
| 53 | |
| 54 | x, y = build_dataset() |
| 55 | test_x = xgb.DMatrix(x) |
| 56 | res = bst.predict(test_x) |
| 57 | assert len(res) == 10 |
| 58 | |
| 59 | def test_training_on_cpu_only_env(self) -> None: |
| 60 | assert os.environ["CUDA_VISIBLE_DEVICES"] == "-1" |
nothing calls this directly
no test coverage detected