I created a model with an embedding layer on CPU, an LSTM on gpu:0 and a softmax layer on gpu:1.
I then created a training op like so:
optim = tf.train.GradientDescentOptimizer(LEARNING_RATE)
train_op = optim.minimize(final_loss)
Training fails with an Out of Memory error. Enabling device placement logging reveals that the gradient for the softmax has been placed on gpu:0 rather than gpu:1. How can get TensorFlow to place it correctly?