assert fails on requesting gpus when not available

5a1d2722 · Yaroslav Bulatov · GitHub · aeb8cfcb · 5a1d2722
Unverified Commit 5a1d2722 authored Jan 10, 2018 by Yaroslav Bulatov Committed by GitHub Jan 10, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

tutorials/image/cifar10_estimator/cifar10_main.py tutorials/image/cifar10_estimator/cifar10_main.py +2 -0

No files found.
--- a/tutorials/image/cifar10_estimator/cifar10_main.py
+++ b/tutorials/image/cifar10_estimator/cifar10_main.py
@@ -503,6 +503,8 @@ if __name__ == '__main__':
      help='Epsilon for batch norm.')
  args = parser.parse_args()

+  if args.num_gpus>0:
+    assert tf.test.is_gpu_available(), "Requested GPUs but none found."
  if args.num_gpus < 0:
    raise ValueError(
        'Invalid GPU count: \"--num-gpus\" must be 0 or a positive integer.')