CODEWITHSUNDEEP
tensorflow
killajoule
killajoule

Reputation: 3832

Tensorflow, how to ensure that everything runs purely in the GPU?

My understanding is that tensorflow will take a computation graph and run it on both the gpu and cpu to max usage. How can I ensure that my code runs on the gpu only?

Upvotes: 0

Views: 98

Answers (1)

mrry
mrry

Reputation: 126154

You can wrap all of your graph-building code in a tf.device() block, as follows:

with tf.device("/gpu:0"):
    # [Build graph in here.]

With this annotation, the placement algorithm will place all operations on the GPU device, and fail if any operations do not have a GPU implementation. (Without this annotation, as you mentioned, it will fall back to placing operations without a GPU implementation on the CPU device.)

Upvotes: 3

Related Questions