Merge pull request tensorflow#3141 from yaroslavvb/patch-1

assert fails on requesting gpus when not available
gipark2001 · Jan 24, 2018 · 1f34fca · 1f34fca
2 parents 5fa2a4e + 6b9d5fb
commit 1f34fca
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tutorials/image/cifar10_estimator/cifar10_main.py b/tutorials/image/cifar10_estimator/cifar10_main.py
@@ -503,6 +503,8 @@ def main(job_dir, data_dir, num_gpus, variable_strategy,
       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.')