Commit 61af56d9 authored by Yuxin Wu's avatar Yuxin Wu

[MaskRCNN] allow inference with MKL support (#1024)

parent 99565998
......@@ -421,7 +421,9 @@ if __name__ == '__main__':
DetectionDataset() # initialize the config with information from our dataset
if args.visualize or args.evaluate or args.predict:
assert tf.test.is_gpu_available()
if not tf.test.is_gpu_available():
from tensorflow.python.framework import test_util
assert test_util.IsMklEnabled(), "Inference requires either GPU support or MKL support!"
assert args.load
finalize_configs(is_training=False)
......
......@@ -44,7 +44,7 @@ def get_num_gpu():
return ret
env = os.environ.get('CUDA_VISIBLE_DEVICES', None)
if env is not None:
if env:
return warn_return(len(env.split(',')), "Found non-empty CUDA_VISIBLE_DEVICES. ")
output, code = subproc_call("nvidia-smi -L", timeout=5)
if code == 0:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment