From 2aa2ff127777c80d38ea05e1cc552d6bba1bd3cd Mon Sep 17 00:00:00 2001 From: Jules Laplace Date: Sat, 14 Apr 2018 19:17:53 +0200 Subject: nsatf at gpu --- python/nsatf.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/nsatf.py b/python/nsatf.py index 6d54765..ce63669 100644 --- a/python/nsatf.py +++ b/python/nsatf.py @@ -29,10 +29,10 @@ else: device_ids = [device.name for device in device_lib.list_local_devices()] -if '/gpu:0' in device_ids: - DEVICE = '/gpu:0' +if '/device:GPU:0' in device_ids: + DEVICE = '/device:GPU:0' else: - DEVICE = '/cpu:0' + DEVICE = '/device:CPU:0' print DEVICE -- cgit v1.2.3-70-g09d2