Using tensorflow-gpu 2.0.0rc0. I want to choose whether it uses the GPU or the CPU.


I've seen some suggestions elsewhere, but they are old and do not apply very well to newer TF versions. What worked for me was this:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

When that variable is defined and equal to -1, TF uses the CPU even when a CUDA GPU is available.

| improve this answer | |
  • 1
    $\begingroup$ I actually find it more convenient to set the environmental variable from outside the script. Sometimes I might forget it in a script and when I import something from that script it automatically runs it and I have no GPU. $\endgroup$ – Djib2011 Sep 7 '19 at 22:39
  • $\begingroup$ At 2.1 (or possibly before) up to nightly, set that environment variable to an empty string to disable GPUs $\endgroup$ – Robert Lugg May 21 at 23:09
  • $\begingroup$ Thanks you so much for this! $\endgroup$ – AnnieLeKatsu Jul 25 at 9:52

For TF2:

    # Disable all GPUS
    tf.config.set_visible_devices([], 'GPU')
    visible_devices = tf.config.get_visible_devices()
    for device in visible_devices:
        assert device.device_type != 'GPU'
    # Invalid device or cannot modify virtual devices once initialized.
| improve this answer | |

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.