check test mode in {cpu, gpu}
(mode="cpu", gpu_id=0)
| 567 | |
| 568 | |
| 569 | def check_test_mode(mode="cpu", gpu_id=0): |
| 570 | """ |
| 571 | check test mode in {cpu, gpu} |
| 572 | """ |
| 573 | if mode == "gpu": |
| 574 | logger.info("using GPU_ID %d", gpu_id) |
| 575 | |
| 576 | elif mode == "cpu": |
| 577 | logger.info("using CPU") |
| 578 | |
| 579 | else: |
| 580 | logger.error( |
| 581 | "Unrecognized argument:%s, 'cpu' or 'gpu' is desired.", mode |
| 582 | ) |
| 583 | log_exit("Invalid arguments") |
| 584 | |
| 585 | return mode |
| 586 | |
| 587 | |
| 588 | def get_test_capacity(run_on_device="cpu"): |
no test coverage detected