diff options
Diffstat (limited to 'python')
| -rw-r--r-- | python/nsatf.py | 6 |
1 files 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 |