with tf.device(DEVICE):
で囲ってやってる