diff options
| author | Jules Laplace <julescarbon@gmail.com> | 2018-04-14 19:17:53 +0200 |
|---|---|---|
| committer | Jules Laplace <julescarbon@gmail.com> | 2018-04-14 19:17:53 +0200 |
| commit | 2aa2ff127777c80d38ea05e1cc552d6bba1bd3cd (patch) | |
| tree | 1b0723bfb94464575cb21295f9b2f5b8c9d9cc6e | |
| parent | 970ad87d53575c80a7dbf751177510e46a3e14dc (diff) | |
nsatf at gpu
| -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 |