How to conditionally assign values to tensor [masking for loss function]?

I want to create a L2 loss function that ignores values (=> pixels) where the label has the value 0. The tensor batch[1] contains the labels while output is a tensor for the net output, both have a shape of (None,300,300,1).

labels_mask = tf.identity(batch[1])
labels_mask[labels_mask > 0] = 1
loss = tf.reduce_sum(tf.square((output-batch[1])*labels_mask))/tf.reduce_sum(labels_mask)

My current code yields to TypeError: 'Tensor' object does not support item assignment (on the second line). What's the tensorflow-way to do this? I also tried to normalize the loss with tf.reduce_sum(labels_mask), which I hope works like this.

If you wanted to write it that way, you would have to use Tensorflow's scatter method for assignment. Unfortunately, tensorflow doesn't really support boolean indexing either (the new boolean_select makes it possible, but annoying). It would be tricky to write and difficult to read.

You have two options that are less annoying:

1. Use labels_mask > 0 as a boolean mask and use Tensorflow's recent boolean_mask function. Maybe this is the more tensorflow way, because it invokes arbitrarily specific functions.
2. Cast labels_mask > 0 to float: tf.cast(labels_mask > 0, tf.float32). Then, you can use it the way you wanted to in the final line of your code.