import tensorflow as tf


config.gpu_options.allow_growth = True
sess_grow = tf.Session(config=config)

# Also, we can limit the size of GPU memory used, with the following option
config.gpu_options.per_process_gpu_memory_fraction = 0.4
sess_limited = tf.Session(config=config)


# How to set placements on multiple devices.
# Here, assume we have three devies CPU:0, GPU:0, and GPU:1
if tf.test.is_built_with_cuda():
    with tf.device('/cpu:0'):
        a = tf.constant([1.0, 3.0, 5.0], shape=[1, 3])
        b = tf.constant([2.0, 4.0, 6.0], shape=[3, 1])
        
        with tf.device('/gpu:1'):
            c = tf.matmul(a,b)
            c = tf.reshape(c, [-1])
        
        with tf.device('/gpu:2'):
            d = tf.matmul(b,a)
            flat_d = tf.reshape(d, [-1])
        
        combined = tf.mul(c, flat_d)
    print(sess.run(combined))