Реализация внимания после слоя CONV и идет не так с большим количеством значений нан

Question

Что-то не так с моей реализацией внимания, потому что я получаю много значений nan.

И мое намерение состоит в том, чтобы приложить внимание перед подачей в слой FC после слоя max-pool CNN.

• Форма ввода -> (10 X 500) <- вывод CNN 2D, а НЕ 3D из RNN

• Форма вывода -> (10 X 500) с учетом внимания (10X1) и применяется

  # attention layer
  # Trainable parameters
  with tf.variable_scope('attention') as scope:
      attention_size = 100
      w_omega = _variable_on_cpu(name='w_omega',
                                 shape=[500,100]
                                 ,initializer=tf.truncated_normal_initializer(stddev=0.01))
      b_omega = _variable_on_cpu(name='b_omega',
                                 shape=[100]
                                 ,initializer=tf.truncated_normal_initializer(stddev=0.01))
      u_omega = _variable_on_cpu(name='u_omega',
                                 shape=[100]
                                 ,initializer=tf.truncated_normal_initializer(stddev=0.01))
      v = tf.tanh(tf.einsum("ab,bc->ac",pool_dropout, w_omega) + b_omega)
      vu = tf.einsum("ab,b->a",v, u_omega)  # (B,T) shape
      alphas = tf.nn.softmax(vu, name='alphas')
      pool_dropout = pool_dropout * tf.expand_dims(alphas,-1)