5fb4404583ecffdee8cf98e4ae9e0a91d479f071,niftynet/layer/loss_regression.py,LossFunction,layer_op,#LossFunction#Any#Any#Any#,26

Before Change

            for ind, pred in enumerate(prediction):
                // go through each scale

                loss_batch = []
                for b_ind, pred_b in enumerate(tf.unstack(pred, axis=0)):
                    // go through each image in a batch

                    pred_b = tf.reshape(pred_b, [-1])

After Change

                    return tf.to_float(self._data_loss_func(**loss_params))

                loss_batch = tf.map_fn(
                    fn=_batch_i_loss,
                    elems=tf.range(tf.shape(pred)[0], dtype=tf.int32),
                    dtype=tf.float32,
                    parallel_iterations=1)
                data_loss.append(tf.reduce_mean(loss_batch))
            return tf.reduce_mean(data_loss)

Instances