Skip to content
Snippets Groups Projects
Commit a7eacc51 authored by Tim O'Donnell's avatar Tim O'Donnell
Browse files

keras backend default tweak

parent 40918bde
No related branches found
No related tags found
No related merge requests found
...@@ -34,7 +34,7 @@ def set_keras_backend(backend=None, gpu_device_nums=None): ...@@ -34,7 +34,7 @@ def set_keras_backend(backend=None, gpu_device_nums=None):
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
[str(i) for i in gpu_device_nums]) [str(i) for i in gpu_device_nums])
if backend == "tensorflow-cpu" or not gpu_device_nums: if backend == "tensorflow-cpu" or gpu_device_nums == []:
print("Forcing tensorflow/CPU backend.") print("Forcing tensorflow/CPU backend.")
os.environ["CUDA_VISIBLE_DEVICES"] = "" os.environ["CUDA_VISIBLE_DEVICES"] = ""
device_count = {'CPU': 1, 'GPU': 0} device_count = {'CPU': 1, 'GPU': 0}
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment