make PerImageStandardization explicitly run on CPU
@@ -549,7 +549,12 @@ class PerImageStandardize(Layer):
return []
def fprop(self, x, **kwargs):
- return tf.map_fn(lambda ex: tf.image.per_image_standardization(ex), x)
+ # TODO: before adding dataset augmentation, we didn't have to do this.
+ # Why?
+ with tf.device("/CPU:0"):
+ out = tf.map_fn(lambda ex: tf.image.per_image_standardization(ex), x)
+ return out
+
class Dropout(Layer):
