Uses gpu when requested and available.
(use_gpu)
| 1450 | |
| 1451 | @contextlib.contextmanager |
| 1452 | def device(use_gpu): |
| 1453 | """Uses gpu when requested and available.""" |
| 1454 | if use_gpu and is_gpu_available(): |
| 1455 | dev = "/device:GPU:0" |
| 1456 | else: |
| 1457 | dev = "/device:CPU:0" |
| 1458 | with ops.device(dev): |
| 1459 | yield |
| 1460 | |
| 1461 | |
| 1462 | @contextlib.contextmanager |
no test coverage detected