Tensorflow ReLu не работает?

Я написал сверточную сеть в тензорном потоке с relu в качестве функции активации, однако она не обучается (потери постоянны для набора данных eval и train). Для разных функций активации все работает как надо.

Вот код, в котором создается nn:

def _create_nn(self):
    current = tf.layers.conv2d(self.input, 20, 3, activation=self.activation)
    current = tf.layers.max_pooling2d(current, 2, 2)
    current = tf.layers.conv2d(current, 24, 3, activation=self.activation)
    current = tf.layers.conv2d(current, 24, 3, activation=self.activation)
    current = tf.layers.max_pooling2d(current, 2, 2)
    self.descriptor = current = tf.layers.conv2d(current, 32, 5, activation=self.activation)
    if not self.drop_conv:
        current = tf.layers.conv2d(current, self.layer_7_filters_n, 3, activation=self.activation)
    if self.add_conv:
        current = tf.layers.conv2d(current, 48, 2, activation=self.activation)

    self.descriptor = current

    last_conv_output_shape = current.get_shape().as_list()
    self.descr_size = last_conv_output_shape[1] * last_conv_output_shape[2] * last_conv_output_shape[3]

    current = tf.layers.dense(tf.reshape(current, [-1, self.descr_size]), 100, activation=self.activation)
    current = tf.layers.dense(current, 50, activation=self.last_activation)

    return current

Для self.activiation установлено значение tf.nn.relu, а для self.last_activiation установлено значение tf.nn.softmax.

Здесь создаются функция потерь и оптимизатор:

    self._nn = self._create_nn()

    self._loss_function = tf.reduce_sum(tf.squared_difference(self._nn, self.Y), 1)

    optimizer = tf.train.AdamOptimizer()
    self._train_op = optimizer.minimize(self._loss_function)

Я попытался изменить инициализацию переменных, передав tf.random_normal_initializer(0.1, 0.1) в качестве инициализаторов, однако это не привело ни к каким изменениям в функции потерь.

Буду благодарен за помощь в настройке этой нейросети на ReLu.

Изменить: у Leaky ReLu такая же проблема

Изменить: небольшой пример, где мне удалось продублировать ту же ошибку:

x = tf.constant([[3., 211., 123., 78.]])
v = tf.Variable([0.5, 0.5, 0.5, 0.5])
h_d = tf.layers.Dense(4, activation=tf.nn.leaky_relu)
h = h_d(x)
y_d = tf.layers.Dense(4, activation=tf.nn.softmax)
y = y_d(h)
d = tf.constant([[.5, .5, 0, 0]])

Градиенты (вычисленные с помощью tf.gradients) для ядер h_d и y_d и смещения либо равны, либо близки к 0