How to set specific gpu in tensorflow?

后端 未结 5 934
没有蜡笔的小新
没有蜡笔的小新 2020-11-29 21:45

I want to specify the gpu to run my process. And I set it as follows:

import tensorflow as tf
with tf.device(\'/gpu:0\'):
    a = tf.constant(3.0)
with tf.Se         


        
5条回答
  •  暖寄归人
    2020-11-29 22:25

    I believe that you need to set CUDA_VISIBLE_DEVICES=1. Or which ever GPU you want to use. If you make only one GPU visible, you will refer to it as /gpu:0 in tensorflow regardless of what you set the environment variable to.

    More info on that environment variable: https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/

提交回复
热议问题