I have a loop in TensorFlow that looks like this:
with tf.device("/gpu:1"):
losses = []
for target, output in zip(targets, lstm_outputs):
logits = tf.matmul(W, output) + b
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits, target)
losses.append(loss)
total_loss = tf.add_n(losses)
I am getting an OOM error when allocating the gradients for this layer, since each matrix multiplication is a different operation in the graph taking memory. Is there a way of preventing TensorFlow from allocating all these operations at the same time?